feat(phyai): Support Cosmos3-Edge - #49
chenghuaWang wants to merge 7 commits into
Conversation
- support Cosmos3-Edge ReLU2 MLP and split Q/K normalization - add a fused Triton ReLU2 kernel - align policy preprocessing, FP32 UniPC state, and linear-flow scheduling
Match Edge and Nano preprocessing, state/history conditioning, action postprocessing, and the native linear-flow UniPC defaults. Add a configurable reference-precision policy modeling path while retaining the fused fast path.
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (12)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThis PR adds Triton FP8, NVFP4, and ReLU2 kernels; introduces model quantization, serialization, and CLI tooling; expands FlashInfer and Humming integration; adds diagnostics and autotuning controls; and updates Cosmos3, PI0.5, benchmarking, and validation paths. ChangesKernel and quantization stack
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (7)
phyai-model-optimizer/pyproject.toml (1)
6-12: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider
~=2.11/==2.11.*for the torch pin.
torch==2.11resolves to the exact2.11.0release and excludes future2.11.xpatch fixes. If patch-level updates should be allowed, usetorch~=2.11.0(or==2.11.*).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@phyai-model-optimizer/pyproject.toml` around lines 6 - 12, Update the torch dependency in the dependencies list to allow compatible 2.11 patch releases, using the project’s preferred compatible-release or wildcard constraint instead of the exact torch==2.11 pin.phyai-model-optimizer/src/phyai_model_optimizer/modifiers/gptq.py (1)
139-142: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueGPTQ is scaffolded but disabled — tracked TODO.
GPTQModifier.__init__raisesNotImplementedError, so the constructor body below (lines 144-175) andgptq_solve's callers are unreachable until calibration tests land. This is fine as an explicit gate; the TODO records the follow-up.One latent gap to fix before enabling:
gptq_solvecomputesngroups = K // group_size(line 53) but never validatesK % group_size == 0. A non-divisibleKmaps trailing columns tog == ngroups, indexingscale[:, g:g+1]out of bounds.Want me to open a tracking issue for enabling GPTQ (with the divisibility guard + calibration tests)?
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@phyai-model-optimizer/src/phyai_model_optimizer/modifiers/gptq.py` around lines 139 - 142, Update gptq_solve to validate that K is evenly divisible by group_size before computing or using ngroups, and reject non-divisible inputs with the established validation/error mechanism. Preserve the existing grouped computation for valid dimensions and prevent trailing columns from producing an out-of-bounds scale-group access.phyai/src/phyai/cli/probe.py (2)
1-1: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
EnvVarProbe.parsed_or_defaultis computed but never displayed.
phyai_env()parses/records the effective value for every registeredPHYAI_*field, but both consumers only ever showraw/error, so the parsed representation is dead weight — anddoctor's own module docstring advertises "PHYAI_*parsing" as a check, which this would make visible.
phyai/src/phyai/cli/probe.py#L384-401: keep computingparsed_or_default, or drop it if intentionally unused.phyai/src/phyai/cli/doctor.py#L220-224: appendvar.parsed_or_defaultto the detail string for successfully-parsed active vars.phyai/src/phyai/cli/info.py#L158-162: same — surfacevar.parsed_or_defaultalongsidevar.rawin the table.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@phyai/src/phyai/cli/probe.py` at line 1, Update the successful active-variable display paths in doctor’s environment checks and info’s environment table to include EnvVarProbe.parsed_or_default alongside the existing raw value; keep error handling and probe computation unchanged.
1-1: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNew "why" comments in this cohort omit the required
# note(name)author tag.Several newly added rationale comments explain why but don't follow the mandated format.
phyai/src/phyai/cli/probe.py#L16-18: add# note(<author>)to the pyproject-mirroring rationale comment.phyai/src/phyai/cli/__init__.py#L9-10: add# note(<author>)to the SUPPRESS-default rationale comment.phyai/src/phyai/cli/doctor.py#L62-63: add# note(<author>)to the optional-extension rationale comment.phyai/src/phyai/cli/probe.py#L268-270: add# note(<author>)to the_MAX_SM_TIERrationale comment.As per coding guidelines, "Comments must be self-contained, concise, written in English, explain why rather than what or how, and include the author's name in the form
# note(name)."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@phyai/src/phyai/cli/probe.py` at line 1, Update the four rationale comments identified around the pyproject-mirroring logic, SUPPRESS default, optional-extension handling, and _MAX_SM_TIER in the referenced CLI modules to include an author tag in the exact `# note(name)` format. Keep each comment self-contained, concise, and focused on why the behavior exists.Source: Coding guidelines
phyai/src/phyai/cli/info.py (1)
148-166: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
_env_sectionnever surfaces the parsed value fetched fromprobe.py.See consolidated comment for details (shared root cause with
doctor.py's_sec_env).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@phyai/src/phyai/cli/info.py` around lines 148 - 166, Update _env_section to display each registered environment variable’s parsed value from probe.phyai_env(), while preserving the existing raw-value and error presentation. Use the parsed-value field exposed by the returned variable objects and keep extra-variable handling unchanged.phyai/src/phyai/cli/doctor.py (1)
215-231: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
_sec_envnever surfaces the parsed value it fetches fromprobe.py.See consolidated comment for details (shared root cause with
info.py's_env_section).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@phyai/src/phyai/cli/doctor.py` around lines 215 - 231, Update _sec_env to display each registered environment variable’s parsed value from probe.phyai_env(), alongside its existing name and raw value output; preserve the current error, warning, default, and parse-failure handling.phyai/src/phyai/models/cosmos3/modeling_cosmos3.py (1)
328-347: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winClass docstring is now stale for the dual-K return.
Cosmos3CausalAttention's class docstring (line 233, unchanged) still says the module "returns(out, K, V)", butforwardnow returnsk_for_gen, which can diverge from the K actually used in the self-attention (k_rotated) wheneverk_norm_und_for_genis active. Worth a one-line docstring update to avoid confusing future readers of the return contract.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@phyai/src/phyai/models/cosmos3/modeling_cosmos3.py` around lines 328 - 347, The Cosmos3CausalAttention class docstring must describe that forward returns `(out, k_for_gen, v)`, where the generation K may differ from the self-attention K when `k_norm_und_for_gen` is enabled. Update only the return-contract wording near the class docstring while preserving the existing behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@benchmark/bench_n_batch_ws1_pi05.py`:
- Around line 108-112: Update the inline comment in the input_ids initialization
to follow the repository rule: explain why token ID 2 is used instead of
describing it as an arbitrary non-pad token, and include the author attribution
in the form # note(name).
In `@CLAUDE.md`:
- Around line 55-60: Correct the typos in the “Code Comment conventions”
guidance: change “whywe” to “why we” and “ust delete it” to “just delete it,”
leaving the surrounding guidance unchanged.
In `@phyai/src/phyai/engine_config.py`:
- Around line 386-392: Update the flashinfer_bf16_backend validation in
RuntimeConfig initialization to verify the value is a string before calling
.lower(). Raise the existing configuration ValueError for None or other
non-string values, while preserving case normalization and valid-backend
validation for strings.
- Around line 71-73: Add "cutile" to the _VALID_FLASHINFER_BF16_BACKENDS
allowlist so RuntimeConfig and PHYAI_FLASHINFER_BF16_BACKEND accept the
supported flashinfer.gemm.mm_bf16 backend, while preserving all existing backend
values.
In `@phyai/src/phyai/models/cosmos3/scheduler_ws1_cosmos3_policy.py`:
- Around line 213-223: Both Cosmos3 schedulers incorrectly cast sampler action
state to model dtype before converting back to FP32. In
phyai/src/phyai/models/cosmos3/scheduler_ws1_cosmos3_policy.py lines 213-223 and
phyai/src/phyai/models/cosmos3/scheduler_wn_cosmos3_policy.py lines 163-172,
update action construction to move to dev without dt; in lines 225-246 and
174-195 respectively, move request.cond_action to dev without dt while retaining
the float conversion. Preserve FP32 sampler state and leave only per-step
transformer inputs cast to model dtype.
---
Nitpick comments:
In `@phyai-model-optimizer/pyproject.toml`:
- Around line 6-12: Update the torch dependency in the dependencies list to
allow compatible 2.11 patch releases, using the project’s preferred
compatible-release or wildcard constraint instead of the exact torch==2.11 pin.
In `@phyai-model-optimizer/src/phyai_model_optimizer/modifiers/gptq.py`:
- Around line 139-142: Update gptq_solve to validate that K is evenly divisible
by group_size before computing or using ngroups, and reject non-divisible inputs
with the established validation/error mechanism. Preserve the existing grouped
computation for valid dimensions and prevent trailing columns from producing an
out-of-bounds scale-group access.
In `@phyai/src/phyai/cli/doctor.py`:
- Around line 215-231: Update _sec_env to display each registered environment
variable’s parsed value from probe.phyai_env(), alongside its existing name and
raw value output; preserve the current error, warning, default, and
parse-failure handling.
In `@phyai/src/phyai/cli/info.py`:
- Around line 148-166: Update _env_section to display each registered
environment variable’s parsed value from probe.phyai_env(), while preserving the
existing raw-value and error presentation. Use the parsed-value field exposed by
the returned variable objects and keep extra-variable handling unchanged.
In `@phyai/src/phyai/cli/probe.py`:
- Line 1: Update the successful active-variable display paths in doctor’s
environment checks and info’s environment table to include
EnvVarProbe.parsed_or_default alongside the existing raw value; keep error
handling and probe computation unchanged.
- Line 1: Update the four rationale comments identified around the
pyproject-mirroring logic, SUPPRESS default, optional-extension handling, and
_MAX_SM_TIER in the referenced CLI modules to include an author tag in the exact
`# note(name)` format. Keep each comment self-contained, concise, and focused on
why the behavior exists.
In `@phyai/src/phyai/models/cosmos3/modeling_cosmos3.py`:
- Around line 328-347: The Cosmos3CausalAttention class docstring must describe
that forward returns `(out, k_for_gen, v)`, where the generation K may differ
from the self-attention K when `k_norm_und_for_gen` is enabled. Update only the
return-contract wording near the class docstring while preserving the existing
behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 0aa9aaf7-f66b-44d4-9025-36cccd06cd63
📒 Files selected for processing (101)
.claude/skills/phyai-jetson-kernel-opt/SKILL.md.claude/skills/phyai-kernel-opt/SKILL.md.gitignoreCLAUDE.mdbenchmark/bench_n_batch.pybenchmark/bench_n_batch_ws1_pi05.pybenchmark/pi05/model_flops.pybenchmark/pi05/profile_pi05.pybenchmark/pi05/run.shexamples/cosmos3/run_cosmos3_policy.pyexamples/cosmos3/run_cosmos3_policy_wn.pyexamples/pi05/quantize_fp8.shexamples/pi05/quantize_int4_int8.shexamples/pi05/quantize_mxfp4_fp8.shphyai-kernel/benchmark/bench_relu2.pyphyai-kernel/phyai_kernel/__init__.pyphyai-kernel/phyai_kernel/triton/__init__.pyphyai-kernel/phyai_kernel/triton/fp8_quant.pyphyai-kernel/phyai_kernel/triton/nvfp4.pyphyai-kernel/phyai_kernel/triton/relu2.pyphyai-kernel/tests/test_fp8_gelu_tanh_quant.pyphyai-kernel/tests/test_relu2.pyphyai-model-optimizer/pyproject.tomlphyai-model-optimizer/src/phyai_model_optimizer/__init__.pyphyai-model-optimizer/src/phyai_model_optimizer/cli.pyphyai-model-optimizer/src/phyai_model_optimizer/compat/__init__.pyphyai-model-optimizer/src/phyai_model_optimizer/compat/ct.pyphyai-model-optimizer/src/phyai_model_optimizer/compat/phyai_model.pyphyai-model-optimizer/src/phyai_model_optimizer/entrypoints.pyphyai-model-optimizer/src/phyai_model_optimizer/modifiers/__init__.pyphyai-model-optimizer/src/phyai_model_optimizer/modifiers/awq.pyphyai-model-optimizer/src/phyai_model_optimizer/modifiers/base.pyphyai-model-optimizer/src/phyai_model_optimizer/modifiers/gptq.pyphyai-model-optimizer/src/phyai_model_optimizer/modifiers/rtn.pyphyai-model-optimizer/src/phyai_model_optimizer/modifiers/smoothquant.pyphyai-model-optimizer/src/phyai_model_optimizer/observers/__init__.pyphyai-model-optimizer/src/phyai_model_optimizer/observers/base.pyphyai-model-optimizer/src/phyai_model_optimizer/observers/hessian.pyphyai-model-optimizer/src/phyai_model_optimizer/observers/minmax.pyphyai-model-optimizer/src/phyai_model_optimizer/orchestrator.pyphyai-model-optimizer/src/phyai_model_optimizer/pipelines/__init__.pyphyai-model-optimizer/src/phyai_model_optimizer/pipelines/base.pyphyai-model-optimizer/src/phyai_model_optimizer/pipelines/datafree.pyphyai-model-optimizer/src/phyai_model_optimizer/pipelines/sequential.pyphyai-model-optimizer/src/phyai_model_optimizer/quant_math.pyphyai-model-optimizer/src/phyai_model_optimizer/recipes.pyphyai-model-optimizer/src/phyai_model_optimizer/serialize.pyphyai-utils-tools/src/phyai_utils_tools/models/cosmos3/processor_cosmos3.pyphyai-utils-tools/src/phyai_utils_tools/models/cosmos3/steps_cosmos3.pyphyai/pyproject.tomlphyai/src/phyai/cli/__init__.pyphyai/src/phyai/cli/__main__.pyphyai/src/phyai/cli/doctor.pyphyai/src/phyai/cli/info.pyphyai/src/phyai/cli/probe.pyphyai/src/phyai/cli/ui.pyphyai/src/phyai/engine.pyphyai/src/phyai/engine_config.pyphyai/src/phyai/env.pyphyai/src/phyai/layers/attention/attention/backends/flashinfer.pyphyai/src/phyai/layers/linear/__init__.pyphyai/src/phyai/layers/linear/backend.pyphyai/src/phyai/layers/linear/backends/__init__.pyphyai/src/phyai/layers/linear/backends/flashinfer.pyphyai/src/phyai/layers/linear/backends/humming.pyphyai/src/phyai/layers/linear/backends/torch.pyphyai/src/phyai/layers/linear/dispatch.pyphyai/src/phyai/layers/linear/layers.pyphyai/src/phyai/layers/linear/registry.pyphyai/src/phyai/layers/mlp/dense_mlp.pyphyai/src/phyai/layers/quant/__init__.pyphyai/src/phyai/layers/quant/fp8.pyphyai/src/phyai/layers/quant/humming.pyphyai/src/phyai/layers/quant/importers/compressed_tensors.pyphyai/src/phyai/layers/quant/importers/fp8.pyphyai/src/phyai/layers/quant/importers/modelopt.pyphyai/src/phyai/layers/quant/materialize.pyphyai/src/phyai/layers/quant/nvfp4.pyphyai/src/phyai/layers/quant/plan.pyphyai/src/phyai/layers/quant/scheme.pyphyai/src/phyai/layers/vocab_embedding/layers.pyphyai/src/phyai/models/cosmos3/configuration_cosmos3.pyphyai/src/phyai/models/cosmos3/main_cosmos3_policy.pyphyai/src/phyai/models/cosmos3/main_cosmos3_policy_wn.pyphyai/src/phyai/models/cosmos3/modeling_cosmos3.pyphyai/src/phyai/models/cosmos3/sampler_unipc.pyphyai/src/phyai/models/cosmos3/scheduler_wn_cosmos3_policy.pyphyai/src/phyai/models/cosmos3/scheduler_ws1_cosmos3_policy.pyphyai/src/phyai/models/pi05/model_runner_pi05.pyphyai/src/phyai/models/pi05/modeling_pi05.pyphyai/src/phyai/models/pi05/scheduler_ws1_pi05.pyphyai/src/phyai/utils/humming.pyphyai/src/phyai/weights/loader.pyphyai/src/phyai/weights/shards.pyphyai/tests/layers/linear/test_kernel_flashinfer.pyphyai/tests/layers/linear/test_layers.pyphyai/tests/layers/mlp/test_dense_mlp.pyphyai/tests/layers/quant/test_humming.pyphyai/tests/layers/quant/test_humming_cuda.pyphyai/tests/layers/quant/test_scale_shards.pythird_party/mirage
| input_ids = torch.zeros( | ||
| batch_size, plugin_cfg.tokenizer_max_length, dtype=torch.int64, device=device | ||
| ) | ||
| input_ids[:, 0] = 2 # any non-pad token id | ||
| lang_lens = torch.ones(batch_size, dtype=torch.int64, device=device) | ||
| input_ids[:, :lang_len] = 2 # any non-pad token id | ||
| lang_lens = torch.full((batch_size,), lang_len, dtype=torch.int64, device=device) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Make the inline comment conform to the repository comment rule.
Use the required attribution and state the rationale.
- input_ids[:, :lang_len] = 2 # any non-pad token id
+ # note(chenghua): A fixed non-pad ID keeps generated prompts deterministic.
+ input_ids[:, :lang_len] = 2As per coding guidelines, “Comments must be self-contained, concise, written in English, explain why rather than what or how, and include the author's name in the form # note(name).”
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| input_ids = torch.zeros( | |
| batch_size, plugin_cfg.tokenizer_max_length, dtype=torch.int64, device=device | |
| ) | |
| input_ids[:, 0] = 2 # any non-pad token id | |
| lang_lens = torch.ones(batch_size, dtype=torch.int64, device=device) | |
| input_ids[:, :lang_len] = 2 # any non-pad token id | |
| lang_lens = torch.full((batch_size,), lang_len, dtype=torch.int64, device=device) | |
| input_ids = torch.zeros( | |
| batch_size, plugin_cfg.tokenizer_max_length, dtype=torch.int64, device=device | |
| ) | |
| # note(chenghua): A fixed non-pad ID keeps generated prompts deterministic. | |
| input_ids[:, :lang_len] = 2 | |
| lang_lens = torch.full((batch_size,), lang_len, dtype=torch.int64, device=device) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@benchmark/bench_n_batch_ws1_pi05.py` around lines 108 - 112, Update the
inline comment in the input_ids initialization to follow the repository rule:
explain why token ID 2 is used instead of describing it as an arbitrary non-pad
token, and include the author attribution in the form # note(name).
Source: Coding guidelines
| ## Code Comment conventions | ||
|
|
||
| All comments should be self-contained. Do not explain how you did something or explain what this code block did; just explain whywe did this. And, also adds your name for all of the comments, like # note(foo). If you don't know user's name, just ask them. | ||
|
|
||
| Comments should be concise. If not explain why, ust delete it. In this sense, most of the AI comments should be removed. | ||
|
|
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Typos in the new comment-convention guidance.
"explain whywe did this" and "If not explain why, ust delete it" — should read "why we" and "just delete it".
✏️ Proposed fix
-All comments should be self-contained. Do not explain how you did something or explain what this code block did; just explain whywe did this. And, also adds your name for all of the comments, like # note(foo). If you don't know user's name, just ask them.
+All comments should be self-contained. Do not explain how you did something or explain what this code block did; just explain why we did this. And also add your name to all comments, like # note(foo). If you don't know the user's name, just ask them.
-Comments should be concise. If not explain why, ust delete it. In this sense, most of the AI comments should be removed.
+Comments should be concise. If a comment doesn't explain why, just delete it. In this sense, most AI-generated comments should be removed.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| ## Code Comment conventions | |
| All comments should be self-contained. Do not explain how you did something or explain what this code block did; just explain whywe did this. And, also adds your name for all of the comments, like # note(foo). If you don't know user's name, just ask them. | |
| Comments should be concise. If not explain why, ust delete it. In this sense, most of the AI comments should be removed. | |
| All comments should be self-contained. Do not explain how you did something or explain what this code block did; just explain why we did this. And also add your name to all comments, like # note(foo). If you don't know the user's name, just ask them. | |
| Comments should be concise. If a comment doesn't explain why, just delete it. In this sense, most AI-generated comments should be removed. |
🧰 Tools
🪛 LanguageTool
[grammar] ~57-~57: Ensure spelling is correct
Context: ... what this code block did; just explain whywe did this. And, also adds your name for ...
(QB_NEW_EN_ORTHOGRAPHY_ERROR_IDS_1)
[style] ~57-~57: Consider removing “of” to be more concise
Context: ... did this. And, also adds your name for all of the comments, like # note(foo). If you don'...
(ALL_OF_THE)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@CLAUDE.md` around lines 55 - 60, Correct the typos in the “Code Comment
conventions” guidance: change “whywe” to “why we” and “ust delete it” to “just
delete it,” leaving the surrounding guidance unchanged.
Source: Linters/SAST tools
| _VALID_FLASHINFER_BF16_BACKENDS: frozenset[str] = frozenset( | ||
| {"auto", "cudnn", "cutlass", "tgv", "cublaslt", "tinygemm"} | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
python - <<'PY'
import inspect
from importlib.metadata import version
import flashinfer.gemm
assert version("flashinfer-python") == "0.6.14"
print(inspect.signature(flashinfer.gemm.mm_bf16))
PYRepository: mingti-org/phyai
Length of output: 277
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '## flashinfer mentions\n'
rg -n "flashinfer|cutile|cutlass|tinygemm|cublaslt|tgv|cudnn" -S \
phyai pyproject.toml poetry.lock uv.lock requirements* .claude/skills 2>/dev/null || true
printf '\n## engine_config outline\n'
ast-grep outline phyai/src/phyai/engine_config.py --view expanded || true
printf '\n## engine_config relevant slice\n'
sed -n '1,180p' phyai/src/phyai/engine_config.pyRepository: mingti-org/phyai
Length of output: 50372
🌐 Web query:
FlashInfer 0.6.14 mm_bf16 cutile backend release notes
💡 Result:
The FlashInfer v0.6.14 release introduced the cuTile backend for mm_bf16 [1]. The cuTile (cuda.tile) backend provides a pure-Python, persistent-scheduled GEMM implementation that supports per-shape exhaustive autotuning and is designed for operations where M, K, and N dimensions align with specific kernel constraints [2][3]. Key details regarding this backend and the v0.6.14 release include: Backend Functionality: The mm_bf16_cutile implementation serves as a thin wrapper that pins alpha=1.0 and beta=0.0, handling layout adjustments to match the upstream mm_bf16 row-major (M, K) and transposed (K, N) input requirements [2]. It is accessible via the backend="cutile" parameter in flashinfer.gemm.mm_bf16 and requires NVIDIA SM >= 90 (Hopper architecture or later) [3]. Release Context: The v0.6.14 version encountered distribution challenges due to package size limitations on PyPI, which prevented the standard upload of the flashinfer-cubin wheel [4]. Users were instructed to install the flashinfer-cubin package directly from the flashinfer.ai wheel index to resolve dependency issues [4]. Subsequent Optimization: Following the v0.6.14 release, further refinements were made to the cuTile kernel, such as specializing the epilogue for beta == 0.0 to skip C-matrix loads, which improves performance and avoids NaN-poisoning from uninitialized buffers [5]. For installation in environments using v0.6.14, the recommended approach is to use the specific index-url provided by the project: pip install flashinfer-python pip install flashinfer-cubin --index-url https://flashinfer.ai/whl pip install flashinfer-jit-cache --index-url https://flashinfer.ai/whl/cu129 # (or cu130) [4]
Citations:
- 1: https://newreleases.io/project/github/flashinfer-ai/flashinfer/release/v0.6.14
- 2: https://github.com/flashinfer-ai/flashinfer/blob/2b150b39/flashinfer/gemm/kernels/cutile/mm_bf16_cutile.py
- 3: https://docs.flashinfer.ai/generated/flashinfer.gemm.mm_bf16.html
- 4: 0.6.14 release is still blocked by a project wide 10GB limit flashinfer-ai/flashinfer#3808
- 5: fix(cutile): skip C loads in the mm_bf16 beta-zero epilogue flashinfer-ai/flashinfer#3981
Allow the supported cutile backend.
flashinfer-python 0.6.14 exposes backend="cutile" for flashinfer.gemm.mm_bf16, but this allowlist rejects it via both RuntimeConfig and PHYAI_FLASHINFER_BF16_BACKEND. Add "cutile" here.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@phyai/src/phyai/engine_config.py` around lines 71 - 73, Add "cutile" to the
_VALID_FLASHINFER_BF16_BACKENDS allowlist so RuntimeConfig and
PHYAI_FLASHINFER_BF16_BACKEND accept the supported flashinfer.gemm.mm_bf16
backend, while preserving all existing backend values.
| gemm_be = self.flashinfer_bf16_backend.lower() | ||
| if gemm_be not in _VALID_FLASHINFER_BF16_BACKENDS: | ||
| raise ValueError( | ||
| f"RuntimeConfig.flashinfer_bf16_backend={gemm_be!r} must be " | ||
| f"one of {sorted(_VALID_FLASHINFER_BF16_BACKENDS)}." | ||
| ) | ||
| object.__setattr__(self, "flashinfer_bf16_backend", gemm_be) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Validate the backend type before normalizing it.
A direct RuntimeConfig(flashinfer_bf16_backend=None) currently raises AttributeError at .lower() instead of the intended configuration ValueError.
Proposed fix
- gemm_be = self.flashinfer_bf16_backend.lower()
+ if not isinstance(self.flashinfer_bf16_backend, str):
+ raise ValueError(
+ "RuntimeConfig.flashinfer_bf16_backend must be a string, got "
+ f"{self.flashinfer_bf16_backend!r}."
+ )
+ gemm_be = self.flashinfer_bf16_backend.lower()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| gemm_be = self.flashinfer_bf16_backend.lower() | |
| if gemm_be not in _VALID_FLASHINFER_BF16_BACKENDS: | |
| raise ValueError( | |
| f"RuntimeConfig.flashinfer_bf16_backend={gemm_be!r} must be " | |
| f"one of {sorted(_VALID_FLASHINFER_BF16_BACKENDS)}." | |
| ) | |
| object.__setattr__(self, "flashinfer_bf16_backend", gemm_be) | |
| if not isinstance(self.flashinfer_bf16_backend, str): | |
| raise ValueError( | |
| "RuntimeConfig.flashinfer_bf16_backend must be a string, got " | |
| f"{self.flashinfer_bf16_backend!r}." | |
| ) | |
| gemm_be = self.flashinfer_bf16_backend.lower() | |
| if gemm_be not in _VALID_FLASHINFER_BF16_BACKENDS: | |
| raise ValueError( | |
| f"RuntimeConfig.flashinfer_bf16_backend={gemm_be!r} must be " | |
| f"one of {sorted(_VALID_FLASHINFER_BF16_BACKENDS)}." | |
| ) | |
| object.__setattr__(self, "flashinfer_bf16_backend", gemm_be) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@phyai/src/phyai/engine_config.py` around lines 386 - 392, Update the
flashinfer_bf16_backend validation in RuntimeConfig initialization to verify the
value is a string before calling .lower(). Raise the existing configuration
ValueError for None or other non-string values, while preserving case
normalization and valid-backend validation for strings.
| ).to(dev) | ||
| action = ( | ||
| torch.from_numpy( | ||
| np.random.RandomState(seed) | ||
| .standard_normal((batch, chunk, ad)) | ||
| .astype("float32") | ||
| ) | ||
| .to(dev, dt) | ||
| .float() | ||
| ) | ||
| action[:, :, raw:] = 0.0 # zero the pad tail beyond the embodiment's dim |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Action noise/conditioning is silently rounded to model dtype before being "restored" to FP32, in both Cosmos3 policy schedulers.
Both schedulers build action and cond_action via .to(dev, dt).float(), which truncates to model dtype (e.g. bf16) and then upcasts back to float32 — the upcast does not recover the precision lost in the intermediate cast. This directly contradicts the documented invariant ("Sampler state stays FP32; only transformer inputs are cast to model dtype", scheduler_ws1_cosmos3_policy.py line 151) and is inconsistent with how video's initial noise is constructed in both files (.to(dev) only, no dtype cast). The per-step denoise loops in both files correctly cast only local model_video/model_action copies to dt for the transformer call, confirming the sampler state itself should stay untouched in FP32.
phyai/src/phyai/models/cosmos3/scheduler_ws1_cosmos3_policy.py#L213-L223: change.to(dev, dt).float()to.to(dev)for theactionnoise tensor.phyai/src/phyai/models/cosmos3/scheduler_ws1_cosmos3_policy.py#L225-L246: changerequest.cond_action.to(dev, dt).float()torequest.cond_action.to(dev).float().phyai/src/phyai/models/cosmos3/scheduler_wn_cosmos3_policy.py#L163-L172: change.to(dev, dt).float()to.to(dev)for theactionnoise tensor.phyai/src/phyai/models/cosmos3/scheduler_wn_cosmos3_policy.py#L174-L195: changerequest.cond_action.to(dev, dt).float()torequest.cond_action.to(dev).float().
📍 Affects 2 files
phyai/src/phyai/models/cosmos3/scheduler_ws1_cosmos3_policy.py#L213-L223(this comment)phyai/src/phyai/models/cosmos3/scheduler_ws1_cosmos3_policy.py#L225-L246phyai/src/phyai/models/cosmos3/scheduler_wn_cosmos3_policy.py#L163-L172phyai/src/phyai/models/cosmos3/scheduler_wn_cosmos3_policy.py#L174-L195
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@phyai/src/phyai/models/cosmos3/scheduler_ws1_cosmos3_policy.py` around lines
213 - 223, Both Cosmos3 schedulers incorrectly cast sampler action state to
model dtype before converting back to FP32. In
phyai/src/phyai/models/cosmos3/scheduler_ws1_cosmos3_policy.py lines 213-223 and
phyai/src/phyai/models/cosmos3/scheduler_wn_cosmos3_policy.py lines 163-172,
update action construction to move to dev without dt; in lines 225-246 and
174-195 respectively, move request.cond_action to dev without dt while retaining
the float conversion. Preserve FP32 sampler state and leave only per-step
transformer inputs cast to model dtype.
1 participant
Summary by CodeRabbit
phyaidiagnostics and environment information commands.phyai-optimizequantization workflows (oneshot and data-free PTQ) with support for integer, FP8, and FP4 formats.relu2kernel.