feat(phyai): humming quant kernels - #48
chenghuaWang wants to merge 4 commits into
Conversation
📝 WalkthroughWalkthroughThis PR adds a new phyai-model-optimizer PTQ package (quant math, modifiers, observers, pipelines, orchestrator, CLI, serialization), new Triton FP8/NVFP4 kernels in phyai-kernel, FlashInfer/Humming quantization backend integration in phyai (linear layers, quant specs, materialize routing), staged CUDA-graph execution for pi0.5 model runner/scheduler, a new phyai doctor/info CLI, and supporting benchmark/config/docs updates. ChangesPTQ toolkit, kernels, and phyai runtime integration
Estimated code review effort: 5 (Critical) | ~150 minutes Sequence Diagram(s)sequenceDiagram
participant Scheduler as PI05WS1Scheduler
participant LLMRunner as PI05LLMRunner
participant PaliGemma as PaliGemmaLanguageModel
participant ExpertRunner as PI05ExpertRunner
Scheduler->>Scheduler: _cached_lang_embeddings / _stage_lang_prefix
Scheduler->>LLMRunner: stage_graph_layout(position_ids, write_indices)
alt staged graph available
Scheduler->>LLMRunner: replay_staged(n_per_sample)
else eager
Scheduler->>LLMRunner: forward(hidden_states)
end
LLMRunner->>PaliGemma: write_prefix_kv(inputs_embeds, ...)
PaliGemma->>PaliGemma: write_kv (final layer only)
Scheduler->>ExpertRunner: noise_input_buffer / replay_staged
ExpertRunner-->>Scheduler: action output
sequenceDiagram
participant CLI as phyai-optimize CLI
participant Entrypoint as model_free_ptq
participant Orchestrator as run_oneshot/resolve_targets
participant Pipeline as SequentialCalibrationPipeline
participant Serialize as serialize.save_state_checkpoint
CLI->>Entrypoint: quantize(checkpoint, modifiers)
Entrypoint->>Orchestrator: resolve_targets(model, modifiers)
Orchestrator->>Pipeline: pipeline.run(targets, dataloader, driver)
Pipeline->>Pipeline: quantize_layer per target (RTN/GPTQ)
Pipeline-->>Orchestrator: quantized targets
Orchestrator->>Serialize: save(model, targets, pack_format)
Serialize-->>CLI: written checkpoint + config.json
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a post-training quantization (PTQ) toolkit (phyai-model-optimizer) supporting data-free and calibration-based flows, adds optimized Triton kernels for FP8 and NVFP4 quantization, and integrates these features into the phyai engine alongside model-level optimizations for the pi05 runner. Feedback on these changes highlights several critical areas for improvement: in fp8_quant.py, clamping amax to a minimum epsilon when COLUMN_MAJOR_SCALES is enabled is necessary to prevent numerical instability, and a block-loop or lower row-width limit should be implemented to avoid compilation failures on large block sizes. In shards.py, potential division-by-zero errors must be guarded against when the replication factor exceeds the parallel size, and .view(-1) should replace .reshape(-1) to prevent silent failures on non-contiguous parameters. Finally, in flashinfer.py, caching or pre-allocating the padded activation tensors is recommended to eliminate allocator overhead on the hot path.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| amax = tl.max(tl.abs(activated), axis=1) | ||
| if COLUMN_MAJOR_SCALES: | ||
| scale_inv = tl.where(amax != 0.0, 448.0 / amax, 1.0) | ||
| scale = 1.0 / scale_inv | ||
| q = tl.clamp(activated * scale_inv[:, None], -448.0, 448.0) | ||
| else: | ||
| scale = tl.maximum(amax, 1e-12) * (1.0 / 448.0) | ||
| q = tl.clamp(activated / scale[:, None], -448.0, 448.0) |
There was a problem hiding this comment.
In _gelu_tanh_and_mul_fp8_group128_kernel, when COLUMN_MAJOR_SCALES is enabled, amax is not clamped to a minimum epsilon (like 1e-12). If amax is extremely small but non-zero (e.g., due to underflow or very small activations), 448.0 / amax can overflow to inf, leading to NaN values when multiplied by activated (e.g., 0.0 * inf). Clamping amax to 1e-12 as done in the else branch and other quantization kernels prevents this numerical instability.
amax = tl.maximum(tl.max(tl.abs(activated), axis=1), 1e-12)\n if COLUMN_MAJOR_SCALES:\n scale_inv = 448.0 / amax\n scale = 1.0 / scale_inv\n q = tl.clamp(activated * scale_inv[:, None], -448.0, 448.0)\n else:\n scale = amax * (1.0 / 448.0)\n q = tl.clamp(activated / scale[:, None], -448.0, 448.0)| rank = mesh.axis_local_rank(axis) // replicate | ||
| world = mesh.axis_size(axis) // replicate |
There was a problem hiding this comment.
If replicate is greater than the tensor-parallel size (mesh.axis_size(axis)), world will be computed as 0 due to integer division. This will cause a ZeroDivisionError when evaluating full % world on line 160. Using max(1, ...) prevents division by zero and gracefully handles cases where the replication factor exceeds the active parallel size (e.g., in single-GPU testing or specific GQA configurations).
rank = mesh.axis_local_rank(axis) // replicate\n world = max(1, mesh.axis_size(axis) // replicate)| dest_off = d * leg.weight_offset // leg.total_weight | ||
| rank = mesh.axis_local_rank(leg.axis) // leg.replicate |
There was a problem hiding this comment.
If leg.replicate is greater than the tensor-parallel size (mesh.axis_size(leg.axis)), world will be computed as 0 due to integer division. This will cause a ZeroDivisionError when evaluating full % world on line 211. Using max(1, ...) prevents division by zero and gracefully handles cases where the replication factor exceeds the active parallel size.
rank = mesh.axis_local_rank(leg.axis) // leg.replicate\n world = max(1, mesh.axis_size(leg.axis) // leg.replicate)| f"shape {tuple(loaded.shape)!r}" | ||
| ) | ||
| param.data.reshape(-1)[slots[shard_id]].copy_(loaded.reshape(())) | ||
|
|
There was a problem hiding this comment.
Using .reshape(-1) on param.data can return a copy instead of a view if the tensor is non-contiguous, which would cause the subsequent .copy_() to silently fail to modify the original parameter in-place. Using .view(-1) is safer because it guarantees a view is returned (or raises an error if the tensor is non-contiguous, preventing silent failures).
| param.data.view(-1)[slots[shard_id]].copy_(loaded.reshape(())) |
| act_x_padded = torch.empty( | ||
| (padded_m, K), dtype=act_x.dtype, device=act_x.device | ||
| ) |
There was a problem hiding this comment.
Allocating act_x_padded dynamically on the hot path of _block_fp8_from_activation using torch.empty can introduce allocator overhead in eager mode. Consider caching these padded tensors on the kernel instance or pre-allocating them to avoid overhead on the critical path.
| if K > _MAX_ROW_WIDTH: | ||
| raise ValueError( | ||
| f"per-token FP8 quantization supports K <= {_MAX_ROW_WIDTH}, got {K}" | ||
| ) |
There was a problem hiding this comment.
If K is close to _MAX_ROW_WIDTH (65536), block_size = triton.next_power_of_2(K) will be extremely large (up to 65536). Triton kernels with such large BLOCK_SIZE will fail to compile or run due to hardware limitations on block size and register/shared memory limits. Consider implementing a block-loop inside _fp8_token_quant_kernel to process K in smaller chunks (e.g., BLOCK_SIZE = 1024 or 2048) or lowering _MAX_ROW_WIDTH to a safe maximum (e.g., 4096 or 8192) if larger dimensions are not expected.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Nitpick comments (6)
examples/pi05/quantize_fp8.sh (1)
15-15: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winKeep the target regex aligned with the documented gate/up/down scope.
Each pattern ends at
.mlp., which may select every descendant under an MLP if the resolver performs regex search matching. Either encode the projection names explicitly or verify thatresolve_targetslimits matches to the intended linear modules.
examples/pi05/quantize_fp8.sh#L15-L15: narrow the suffix or verify resolver filtering.examples/pi05/quantize_int4_int8.sh#L15-L15: narrow the suffix or verify resolver filtering.examples/pi05/quantize_mxfp4_fp8.sh#L15-L15: narrow the suffix or verify resolver filtering.🤖 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 `@examples/pi05/quantize_fp8.sh` at line 15, Ensure the target regexes only match the documented gate/up/down linear projections rather than every descendant under an MLP. Update the patterns at examples/pi05/quantize_fp8.sh:15-15, examples/pi05/quantize_int4_int8.sh:15-15, and examples/pi05/quantize_mxfp4_fp8.sh:15-15 to explicitly narrow the projection suffixes, or verify and rely on resolve_targets filtering if it already enforces the intended scope.phyai-model-optimizer/src/phyai_model_optimizer/serialize.py (1)
294-330: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffConsider extracting the shared pack-format validation rules. Both functions encode the same constraints (num_bits∈{3,6} → humming-only, E5M2 not representable in CT float8, FP4 → compressed-tensors-only) but against different inputs, so a rule change in one spot can silently drift from the other.
phyai-model-optimizer/src/phyai_model_optimizer/serialize.py#L294-L330: keep the target-facing entry but route the dtype-rule checks through a shared helper taking an iterable ofWeightQuant.phyai-model-optimizer/src/phyai_model_optimizer/entrypoints.py#L24-L58: call the same shared helper (mapping modifiers →weight_quant()) instead of re-implementing the rules.🤖 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/serialize.py` around lines 294 - 330, Extract the shared dtype-rule validation from _validate_pack_format into a helper accepting an iterable of WeightQuant, while keeping _validate_pack_format responsible for target iteration and pack-format validation; update serialize.py lines 294-330 to use it. Replace the duplicated rules in entrypoints.py lines 24-58 with the same helper, passing each modifier’s weight_quant(), so num_bits 3/6, E5M2, and FP4 constraints remain consistent.phyai-kernel/tests/test_fp8_gelu_tanh_quant.py (1)
1-14: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winUntested new quant kernels: only the fused GeGLU-tanh path has coverage. This cohort adds six new production numerical kernels (five FP8 activation/weight helpers plus the NVFP4 output-scaling epilogue), but only
gelu_tanh_and_mul_fp8_group128is exercised by a test. These are correctness-critical paths feeding weight/activation quantization for inference (e.g. consumed byphyai/src/phyai/layers/quant/humming.py'sfp8_quantize_weight_per_blockcall and FlashInfer'snvfp4_scale_outputcall).
phyai-kernel/tests/test_fp8_gelu_tanh_quant.py#L1-L14: add tests forfp8_quantize_per_tensor,fp8_quantize_per_tensor_with_scale,fp8_quantize_per_token,fp8_quantize_weight_per_block, andfp8_requantize_with_scale_ratio(e.g. against a naive PyTorch reference).phyai-kernel/phyai_kernel/triton/nvfp4.py#L79-L102: add a new test file validatingnvfp4_scale_output's per-row scaling, optional bias fusion, and input-validation error paths.🤖 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-kernel/tests/test_fp8_gelu_tanh_quant.py` around lines 1 - 14, Add coverage in phyai-kernel/tests/test_fp8_gelu_tanh_quant.py for fp8_quantize_per_tensor, fp8_quantize_per_tensor_with_scale, fp8_quantize_per_token, fp8_quantize_weight_per_block, and fp8_requantize_with_scale_ratio using naive PyTorch references. Add a separate test file for nvfp4_scale_output in phyai-kernel/phyai_kernel/triton/nvfp4.py, covering per-row scaling, optional bias fusion, and input-validation errors.phyai/src/phyai/layers/quant/importers/modelopt.py (1)
131-144: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated matcher-expansion helper across importers.
_exclude_matchersand_ignored_matchersare identical (glob/name kind detection plus the same fused-name alias map). Extract one shared helper so the fused-name aliases and kind detection stay in sync.
phyai/src/phyai/layers/quant/importers/modelopt.py#L131-L144: replace_exclude_matcherswith a call to the shared helper.phyai/src/phyai/layers/quant/importers/fp8.py#L11-L23: replace_ignored_matcherswith the same shared helper.🤖 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/layers/quant/importers/modelopt.py` around lines 131 - 144, Extract the duplicated matcher-expansion logic from _exclude_matchers in modelopt.py and _ignored_matchers in fp8.py into one shared helper, preserving glob/name detection and the fused-name alias map. Replace both importer-local helpers with calls to the shared helper, updating phyai/src/phyai/layers/quant/importers/modelopt.py lines 131-144 and phyai/src/phyai/layers/quant/importers/fp8.py lines 11-23; both sites require direct changes.benchmark/bench_n_batch_ws1_pi05.py (1)
199-206: 📐 Maintainability & Code Quality | 🔵 TrivialLang-bucket logic duplicates
model_flops.bucket_lang_lenwith a hardcoded set.
lang_buckets/lang_buckethere hardcode(16, 48, 112)plustokenizer_max_length, re-implementing the same concept asmf.bucket_lang_len(already used for this purpose inbenchmark/pi05/profile_pi05.py). If the canonical bucket set changes there, this script's reportedlang_bucketextras silently drift out of sync.♻️ Reuse the shared bucket helper
- lang_buckets = sorted( - {b for b in (16, 48, 112) if b < plugin_cfg.tokenizer_max_length} - | {plugin_cfg.tokenizer_max_length} - ) - lang_bucket = next( - (bucket for bucket in lang_buckets if bucket >= lang_len), - plugin_cfg.tokenizer_max_length, - ) + lang_bucket = mf.bucket_lang_len(lang_len, dims)🤖 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 199 - 206, Replace the hardcoded lang_buckets/lang_bucket calculation with the shared model_flops.bucket_lang_len helper, matching the usage in profile_pi05.py. Pass lang_len and the configured tokenizer_max_length as required, and preserve the resulting lang_bucket value used by the benchmark.phyai/src/phyai/cli/probe.py (1)
16-18: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNew comments across the CLI cohort omit the required author annotation. Every explanatory comment added in this cohort explains the rationale ("why") but none use the
# note(name)tag the guideline requires for Python comments.
phyai/src/phyai/cli/probe.py#L16-L18: tag the pins-mirroring rationale comment, e.g.# note(<author>): Hard pins mirrored from phyai/pyproject.toml....phyai/src/phyai/cli/__init__.py#L9-L10: tag theSUPPRESSrationale comment above_common_flags.phyai/src/phyai/cli/ui.py#L10-L15: tag the brand-mark and status-vocabulary comments.phyai/src/phyai/cli/doctor.py#L62-L63: tag the "optional workspace members" rationale comment above_OPTIONAL_EXT.As per coding guidelines,
**/*.{py,pyi}: "Write comments in English, keep them concise and self-contained, explain why rather than what or how, and add the required author annotation such as# 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` around lines 16 - 18, Annotate each explanatory comment with the required author tag: in phyai/src/phyai/cli/probe.py lines 16-18 tag the hard-pins rationale; in phyai/src/phyai/cli/__init__.py lines 9-10 tag the SUPPRESS rationale above _common_flags; in phyai/src/phyai/cli/ui.py lines 10-15 tag both brand-mark and status-vocabulary comments; and in phyai/src/phyai/cli/doctor.py lines 62-63 tag the optional-workspace-members rationale above _OPTIONAL_EXT. Keep the comments concise, English, and otherwise unchanged.Source: Coding guidelines
🤖 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 `@CLAUDE.md`:
- Around line 55-60: Correct the malformed wording in the “Code Comment
conventions” section of CLAUDE.md: change “whywe” to “why we” and “ust” to
“just,” leaving the surrounding guidance unchanged.
In `@examples/pi05/quantize_fp8.sh`:
- Around line 16-24: Normalize the input and output positional paths before
changing directories, then pass those normalized paths to the quantization
command. Apply this consistently in examples/pi05/quantize_fp8.sh lines 16-24,
examples/pi05/quantize_int4_int8.sh lines 17-25, and
examples/pi05/quantize_mxfp4_fp8.sh lines 17-25, preserving each script’s
existing quantization options.
In `@phyai/src/phyai/cli/probe.py`:
- Around line 170-177: Broaden the exception handling in _run so any Exception
raised by subprocess.run returns None, matching the module’s other probe helpers
and ensuring CLI commands do not crash on permission or other subprocess errors.
Keep the existing command execution and combined stdout/stderr return behavior
unchanged.
In `@phyai/src/phyai/layers/linear/dispatch.py`:
- Around line 71-77: Update the cache-key construction in the dispatch method to
remove exact M while preserving whether M is None: use Mb together with an
explicit M-is-None indicator, alongside the existing spec, dimensions, dtypes,
device capability, and mode fields. Keep can_handle behavior unchanged and
ensure all cache lookups use the revised key consistently.
In `@phyai/src/phyai/layers/quant/nvfp4.py`:
- Around line 40-49: Update flashinfer_nvfp4_e4m3_max so the try/except covers
both loading and invoking current_nvfp4_4over6_config and nvfp4_e4m3_max. Catch
the same AttributeError and ImportError from either the import or runtime call,
returning _FP8_E4M3_AMAX whenever FlashInfer is incompatible.
In `@phyai/src/phyai/models/pi05/scheduler_ws1_pi05.py`:
- Around line 215-220: Update _tensor_version to avoid relying solely on
Tensor._version: retain the current version lookup when supported, but add a
defensive compatibility fallback for missing attributes or alternate failures so
cache invalidation remains safe. Document the fallback’s compatibility purpose
clearly within the helper.
In `@phyai/tests/layers/quant/test_scale_shards.py`:
- Around line 179-184: Update the helper function _layer_with_scale to attach
the provided extra metadata to the parameter’s _quant_attrs attribute instead of
_humming_attrs, so _attach_optional_scales can detect it and the tests exercise
scale_sharded and scale_fused rather than falling back to replicated().
---
Nitpick comments:
In `@benchmark/bench_n_batch_ws1_pi05.py`:
- Around line 199-206: Replace the hardcoded lang_buckets/lang_bucket
calculation with the shared model_flops.bucket_lang_len helper, matching the
usage in profile_pi05.py. Pass lang_len and the configured tokenizer_max_length
as required, and preserve the resulting lang_bucket value used by the benchmark.
In `@examples/pi05/quantize_fp8.sh`:
- Line 15: Ensure the target regexes only match the documented gate/up/down
linear projections rather than every descendant under an MLP. Update the
patterns at examples/pi05/quantize_fp8.sh:15-15,
examples/pi05/quantize_int4_int8.sh:15-15, and
examples/pi05/quantize_mxfp4_fp8.sh:15-15 to explicitly narrow the projection
suffixes, or verify and rely on resolve_targets filtering if it already enforces
the intended scope.
In `@phyai-kernel/tests/test_fp8_gelu_tanh_quant.py`:
- Around line 1-14: Add coverage in
phyai-kernel/tests/test_fp8_gelu_tanh_quant.py for fp8_quantize_per_tensor,
fp8_quantize_per_tensor_with_scale, fp8_quantize_per_token,
fp8_quantize_weight_per_block, and fp8_requantize_with_scale_ratio using naive
PyTorch references. Add a separate test file for nvfp4_scale_output in
phyai-kernel/phyai_kernel/triton/nvfp4.py, covering per-row scaling, optional
bias fusion, and input-validation errors.
In `@phyai-model-optimizer/src/phyai_model_optimizer/serialize.py`:
- Around line 294-330: Extract the shared dtype-rule validation from
_validate_pack_format into a helper accepting an iterable of WeightQuant, while
keeping _validate_pack_format responsible for target iteration and pack-format
validation; update serialize.py lines 294-330 to use it. Replace the duplicated
rules in entrypoints.py lines 24-58 with the same helper, passing each
modifier’s weight_quant(), so num_bits 3/6, E5M2, and FP4 constraints remain
consistent.
In `@phyai/src/phyai/cli/probe.py`:
- Around line 16-18: Annotate each explanatory comment with the required author
tag: in phyai/src/phyai/cli/probe.py lines 16-18 tag the hard-pins rationale; in
phyai/src/phyai/cli/__init__.py lines 9-10 tag the SUPPRESS rationale above
_common_flags; in phyai/src/phyai/cli/ui.py lines 10-15 tag both brand-mark and
status-vocabulary comments; and in phyai/src/phyai/cli/doctor.py lines 62-63 tag
the optional-workspace-members rationale above _OPTIONAL_EXT. Keep the comments
concise, English, and otherwise unchanged.
In `@phyai/src/phyai/layers/quant/importers/modelopt.py`:
- Around line 131-144: Extract the duplicated matcher-expansion logic from
_exclude_matchers in modelopt.py and _ignored_matchers in fp8.py into one shared
helper, preserving glob/name detection and the fused-name alias map. Replace
both importer-local helpers with calls to the shared helper, updating
phyai/src/phyai/layers/quant/importers/modelopt.py lines 131-144 and
phyai/src/phyai/layers/quant/importers/fp8.py lines 11-23; both sites require
direct changes.
🪄 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: 5af3d3e6-bba5-4f3d-9ba4-2d6296676b92
📒 Files selected for processing (88)
.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_wn.pyexamples/pi05/quantize_fp8.shexamples/pi05/quantize_int4_int8.shexamples/pi05/quantize_mxfp4_fp8.shphyai-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/tests/test_fp8_gelu_tanh_quant.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/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/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
💤 Files with no reviewable changes (1)
- examples/cosmos3/run_cosmos3_policy_wn.py
| ## 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
Fix the malformed comment-convention wording.
Correct whywe to why we and ust to just; otherwise agents may misread these repository instructions.
Suggested wording
-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. Explain why the code exists, not how it was implemented or what it does. Add an author annotation such as # note(foo); if the user's name is unknown, ask for it.
-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 does not explain why, just delete it. Remove unnecessary AI-generated commentary.📝 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. | |
| ## Code Comment conventions | |
| All comments should be self-contained. Explain why the code exists, not how it was implemented or what it does. Add an author annotation such as # note(foo); if the user's name is unknown, ask for it. | |
| Comments should be concise. If a comment does not explain why, just delete it. Remove unnecessary AI-generated commentary. |
🧰 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 malformed wording in the “Code
Comment conventions” section of CLAUDE.md: change “whywe” to “why we” and “ust”
to “just,” leaving the surrounding guidance unchanged.
Sources: Coding guidelines, Linters/SAST tools
| cd "${REPO_ROOT}" | ||
| uv run phyai-optimize quantize \ | ||
| --input "$1" \ | ||
| --output "$2" \ | ||
| --weight-dtype fp8_e4m3 \ | ||
| --activation-dtype fp8_e4m3 \ | ||
| --fp8-scheme block-128 \ | ||
| --pack-format compressed-tensors \ | ||
| --targets "${MLP_TARGET}" |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Resolve user paths before changing directories in all quantization launchers.
All three scripts change to the repository root before forwarding positional arguments, so relative paths depend on the caller’s working directory and can fail or target the wrong location.
examples/pi05/quantize_fp8.sh#L16-L24: normalize$1and$2beforecd, then pass the normalized paths.examples/pi05/quantize_int4_int8.sh#L17-L25: apply the same path normalization.examples/pi05/quantize_mxfp4_fp8.sh#L17-L25: apply the same path normalization.
📍 Affects 3 files
examples/pi05/quantize_fp8.sh#L16-L24(this comment)examples/pi05/quantize_int4_int8.sh#L17-L25examples/pi05/quantize_mxfp4_fp8.sh#L17-L25
🤖 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 `@examples/pi05/quantize_fp8.sh` around lines 16 - 24, Normalize the input and
output positional paths before changing directories, then pass those normalized
paths to the quantization command. Apply this consistently in
examples/pi05/quantize_fp8.sh lines 16-24, examples/pi05/quantize_int4_int8.sh
lines 17-25, and examples/pi05/quantize_mxfp4_fp8.sh lines 17-25, preserving
each script’s existing quantization options.
| def _run(cmd: list[str], timeout: float = 8.0) -> str | None: | ||
| try: | ||
| done = subprocess.run( | ||
| cmd, capture_output=True, text=True, timeout=timeout, check=False | ||
| ) | ||
| except (FileNotFoundError, subprocess.TimeoutExpired): | ||
| return None | ||
| return (done.stdout or "") + (done.stderr or "") |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Broaden _run()'s exception handling to match the rest of the module.
Only FileNotFoundError/TimeoutExpired are caught here, but every other probe in this file uses a broad except Exception to guarantee it never crashes the CLI. A PermissionError or other OSError from subprocess.run (restrictive perms on nvcc/git/nvidia-smi in a sandboxed/CI host) would propagate uncaught and crash phyai doctor/phyai info.
🛡️ Proposed fix
- try:
- done = subprocess.run(
- cmd, capture_output=True, text=True, timeout=timeout, check=False
- )
- except (FileNotFoundError, subprocess.TimeoutExpired):
- return None
+ try:
+ done = subprocess.run(
+ cmd, capture_output=True, text=True, timeout=timeout, check=False
+ )
+ except (OSError, subprocess.TimeoutExpired):
+ return None📝 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.
| def _run(cmd: list[str], timeout: float = 8.0) -> str | None: | |
| try: | |
| done = subprocess.run( | |
| cmd, capture_output=True, text=True, timeout=timeout, check=False | |
| ) | |
| except (FileNotFoundError, subprocess.TimeoutExpired): | |
| return None | |
| return (done.stdout or "") + (done.stderr or "") | |
| def _run(cmd: list[str], timeout: float = 8.0) -> str | None: | |
| try: | |
| done = subprocess.run( | |
| cmd, capture_output=True, text=True, timeout=timeout, check=False | |
| ) | |
| except (OSError, subprocess.TimeoutExpired): | |
| return None | |
| return (done.stdout or "") + (done.stderr or "") |
🧰 Tools
🪛 ast-grep (0.44.1)
[error] 171-173: Use of unsanitized data to create processes
Context: subprocess.run(
cmd, capture_output=True, text=True, timeout=timeout, check=False
)
Note: [CWE-78] Improper Neutralization of Special Elements used in an OS Command ('OS Command Injection').
(os-system-unsanitized-data)
[error] 171-173: Command coming from incoming request
Context: subprocess.run(
cmd, capture_output=True, text=True, timeout=timeout, check=False
)
Note: [CWE-78] Improper Neutralization of Special Elements used in an OS Command ('OS Command Injection').
(subprocess-from-request)
🤖 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` around lines 170 - 177, Broaden the exception
handling in _run so any Exception raised by subprocess.run returns None,
matching the module’s other probe helpers and ensuring CLI commands do not crash
on permission or other subprocess errors. Keep the existing command execution
and combined stdout/stderr return behavior unchanged.
| mode = current_mode() | ||
| Mb = _M_bucket(M) | ||
| key = (spec_id, Mb, N, K, in_dtype, out_dtype, self._sm, mode) | ||
| sm = self._sm_by_device.get(device) | ||
| if sm is None: | ||
| sm = sm_arch(device) | ||
| self._sm_by_device[device] = sm | ||
| key = (spec_id, M, Mb, N, K, in_dtype, out_dtype, sm, mode) |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Does any LinearKernel.can_handle branch on the numeric value of probe.M
# (beyond a None-check)? If none do, exact-M keying only bloats the cache.
fd -e py . phyai/src/phyai/layers/linear/backends \
--exec rg -nP -C3 '\bprobe\.M\b'Repository: mingti-org/phyai
Length of output: 795
Cache should not key on exact M
self._cache now creates a separate entry for every distinct token count, but the visible can_handle branches only depend on M is None plus dtype/alignment checks. If the kernel choice does not need the numeric M, key this on M_bucket with an M is None flag instead.
🤖 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/layers/linear/dispatch.py` around lines 71 - 77, Update the
cache-key construction in the dispatch method to remove exact M while preserving
whether M is None: use Mb together with an explicit M-is-None indicator,
alongside the existing spec, dimensions, dtypes, device capability, and mode
fields. Keep can_handle behavior unchanged and ensure all cache lookups use the
revised key consistently.
| def flashinfer_nvfp4_e4m3_max() -> float: | ||
| """Return FlashInfer's active NVFP4 E4M3 limit, including 4-over-6.""" | ||
| try: | ||
| from flashinfer.quantization.nvfp4_quantization_utils import ( | ||
| current_nvfp4_4over6_config, | ||
| nvfp4_e4m3_max, | ||
| ) | ||
| except (AttributeError, ImportError): | ||
| return _FP8_E4M3_AMAX | ||
| return nvfp4_e4m3_max(current_nvfp4_4over6_config()) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
flashinfer 0.6.14 nvfp4_quantization_utils current_nvfp4_4over6_config nvfp4_e4m3_max
💡 Result:
In FlashInfer 0.6.14, the nvfp4_quantization_utils module provides essential configuration, utilities, and environment-driven settings to support NVFP4 quantization, specifically focusing on "4-over-6" (4over6) quantization modes [1][2]. Key components related to your query include: - current_nvfp4_4over6_config: This utility function retrieves the active NVFP4 4-over-6 configuration [1][3]. It inspects environment variables—such as FLASHINFER_NVFP4_4OVER6, FLASHINFER_NVFP4_4OVER6_ERR_MODE, and FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH—to construct an NVFP44Over6Config dataclass, which standardizes quantization parameters used across different kernels [1][2]. - nvfp4_e4m3_max: This function (or property) returns the maximum value for the E4M3 scale factor based on the provided NVFP44Over6Config [3][4]. The 4-over-6 recipe supports different E4M3 bounds (commonly 256 or 448), which are crucial for defining the dynamic range of the quantized tensors [5][6]. These utilities allow the library to maintain a consistent quantization configuration between Python-side helpers and underlying C++/CUDA kernels, ensuring that scaling factors and error-calculation modes are applied correctly during NVFP4 quantization and MoE pipelines [1][2].
Citations:
- 1: Add CuTe DSL NVFP4 quantization with 4over6 FP16 scoring flashinfer-ai/flashinfer#3448
- 2: Support 4over6 nvfp4 for quantizer and fused MoE flashinfer-ai/flashinfer#3264
- 3: https://github.com/flashinfer-ai/flashinfer/blob/2b150b39/benchmarks/routines/quantization.py
- 4: https://github.com/flashinfer-ai/flashinfer/blob/2b150b39/tests/test_helpers/utils_fp4.py
- 5: Implement 4over6 NVFP4 recipe NVIDIA/TransformerEngine#2972
- 6: Optimize NVFP4 4over6 candidate error path NVIDIA/TransformerEngine#3068
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the target file around the helper and any nearby callers.
sed -n '1,120p' phyai/src/phyai/layers/quant/nvfp4.py
printf '\n--- SEARCH: helper usage ---\n'
rg -n "flashinfer_nvfp4_e4m3_max|current_nvfp4_4over6_config|nvfp4_e4m3_max" phyai/src/phyai/layers/quant/nvfp4.pyRepository: mingti-org/phyai
Length of output: 4700
Extend the FlashInfer fallback to cover the runtime call.
phyai/src/phyai/layers/quant/nvfp4.py:40-49 — the try/except only wraps the import. If current_nvfp4_4over6_config() or nvfp4_e4m3_max(...) raises on an incompatible FlashInfer build, this helper still crashes instead of returning _FP8_E4M3_AMAX.
🤖 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/layers/quant/nvfp4.py` around lines 40 - 49, Update
flashinfer_nvfp4_e4m3_max so the try/except covers both loading and invoking
current_nvfp4_4over6_config and nvfp4_e4m3_max. Catch the same AttributeError
and ImportError from either the import or runtime call, returning _FP8_E4M3_AMAX
whenever FlashInfer is incompatible.
| def _tensor_version(tensor: torch.Tensor) -> int | None: | ||
| try: | ||
| return int(tensor._version) | ||
| except RuntimeError: | ||
| return None | ||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
Is torch.Tensor._version a stable public-facing API in PyTorch, or purely internal/undocumented?
💡 Result:
The torch.Tensor._version attribute is not a stable, public-facing API in PyTorch; it is an internal, undocumented implementation detail [1][2][3]. While it is accessible on tensor objects and is frequently used by developers for debugging purposes—such as monitoring when a tensor has been modified by in-place operations—it is not intended for use in production code or dependent logic [1][2][4][3]. Because it is internal, its behavior is subject to change without notice in future PyTorch versions [2][3]. Reliance on this attribute is generally discouraged as it may break or behave unexpectedly in certain contexts, such as when using torch.inference_mode where autograd is disabled and the version counter is not available [2].
Citations:
- 1: Tensor._version isn't updated when wrapped in a Parameter pytorch/pytorch#7386
- 2: https://discuss.pytorch.org/t/how-to-get-the-version-numbers-of-a-modules-parameters/90726
- 3: https://discuss.pytorch.org/t/getting-the-autograd-counter-of-a-tensor-in-pytorch/126345
- 4: https://stackoverflow.com/questions/68326500/getting-the-autograd-counter-of-a-tensor-in-pytorch
Avoid depending on Tensor._version for cache invalidation
tensor._version is an internal PyTorch detail, so this cache gate can silently stop invalidating if the attribute changes or a tensor subclass raises a different exception. Add a defensive fallback here, or wrap the version read in a helper with a clear compatibility note.
🤖 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/pi05/scheduler_ws1_pi05.py` around lines 215 - 220,
Update _tensor_version to avoid relying solely on Tensor._version: retain the
current version lookup when supported, but add a defensive compatibility
fallback for missing attributes or alternate failures so cache invalidation
remains safe. Document the fallback’s compatibility purpose clearly within the
helper.
| def _layer_with_scale(shape, extra): | ||
| layer = nn.Module() | ||
| p = nn.Parameter(torch.zeros(*shape), requires_grad=False) | ||
| p._humming_attrs = extra | ||
| layer.weight_scale = p | ||
| return layer |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm the attribute name the dispatcher actually reads.
rg -nP '_quant_attrs|_humming_attrs|_quant_source_name' phyai/src/phyai/layers/linear/layers.pyRepository: mingti-org/phyai
Length of output: 564
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== test file excerpt =="
sed -n '160,230p' phyai/tests/layers/quant/test_scale_shards.py | cat -n
echo
echo "== loader excerpt =="
sed -n '120,180p' phyai/src/phyai/layers/linear/layers.py | cat -nRepository: mingti-org/phyai
Length of output: 6810
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n "def replicated|def scale_sharded|def scale_fused|weight_loader" phyai/src/phyai/layers/linear/layers.pyRepository: mingti-org/phyai
Length of output: 1862
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1,260p' phyai/src/phyai/layers/linear/layers.py | cat -nRepository: mingti-org/phyai
Length of output: 11257
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n "def replicated|def scale_sharded|def scale_fused|class .*Loader|weight_loader" phyai/src/phyai/weights/shards.py
sed -n '1,260p' phyai/src/phyai/weights/shards.py | cat -nRepository: mingti-org/phyai
Length of output: 12149
Use _quant_attrs in _layer_with_scale
_attach_optional_scales looks for _quant_attrs, so _humming_attrs leaves extra empty and both tests fall back to replicated() instead of exercising scale_sharded / scale_fused.
🐛 Proposed fix
def _layer_with_scale(shape, extra):
layer = nn.Module()
p = nn.Parameter(torch.zeros(*shape), requires_grad=False)
- p._humming_attrs = extra
+ p._quant_attrs = extra
layer.weight_scale = p
return layer📝 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.
| def _layer_with_scale(shape, extra): | |
| layer = nn.Module() | |
| p = nn.Parameter(torch.zeros(*shape), requires_grad=False) | |
| p._humming_attrs = extra | |
| layer.weight_scale = p | |
| return layer | |
| def _layer_with_scale(shape, extra): | |
| layer = nn.Module() | |
| p = nn.Parameter(torch.zeros(*shape), requires_grad=False) | |
| p._quant_attrs = extra | |
| layer.weight_scale = p | |
| return layer |
🤖 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/tests/layers/quant/test_scale_shards.py` around lines 179 - 184, Update
the helper function _layer_with_scale to attach the provided extra metadata to
the parameter’s _quant_attrs attribute instead of _humming_attrs, so
_attach_optional_scales can detect it and the tests exercise scale_sharded and
scale_fused rather than falling back to replicated().
1 participant
Summary by CodeRabbit
New Features
phyaidiagnostics and environment information commands.Documentation
Bug Fixes