Skip to content

[None][chore] Core: MoE backend enum, deepgemm JIT warmup, postproc metrics, cleanup - #19186

Closed
reasonsolo wants to merge 4 commits into
NVIDIA:mainfrom
reasonsolo:user/lizhiz/mb-core
Closed

[None][chore] Core: MoE backend enum, deepgemm JIT warmup, postproc metrics, cleanup#19186
reasonsolo wants to merge 4 commits into
NVIDIA:mainfrom
reasonsolo:user/lizhiz/mb-core

Conversation

@reasonsolo

@reasonsolo reasonsolo commented Sep 15, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Description

Core/shared changes from the Rubin merge-back series.

  • MoeConfig.backend gains "CUTEDSL_FC12" plus the regenerated llm_args_golden_manifest.json. Its consumer, moe_resolution.py → CuteDslFc12FusedMoE, lands in the MoE PR of this series.
  • deep_gemm_jit_warmup_buckets in _torch/utils.py and both call sites in torch_custom_ops.py. The pre-existing deep_gemm_gen_tuning_buckets is untouched.
  • Postproc metric assignment (executor/result.py) and a dynamo-GC fix (executor/worker.py).
  • Small fixes against main's own code: dwdp/setup.py (K3 does use block_sparse_moe/routed_experts) and modeling_kimi_k25.py (_get_vision_tp_mapping ignored cp_size).
  • A CuTe DSL compatibility shim in tensorrt_llm/__init__.py.

Scope: 18 of 30 candidate files

Twelve were dropped after verification. The notable ones:

  • compilation/patterns/residual_add_norm.py and compilation/utils.py — main already has the fix. compilation/patterns/__init__.py:24 defines _make_pattern_example_inputs with the same Torch 2.13+ _trace_args_for_initial_trace rationale, and derives arity from signature(search_fn) rather than hardcoding dummies. Strictly better than the Rubin version; nothing to port.
  • Three of four candidate config fields were dropped, each verified unbacked:
    • uses_spcompress would raise TypeError — SkipSoftmaxParams is a frozen dataclass whose only fields are algorithm and scheduler, on main and on the Rubin branch.
    • use_lamport_sync has no reader anywhere on either branch.
    • enable_ugpu has no consumer on either branch.
  • test_py_executor.py — six new cases call PyExecutor._terminate_attention_dp_dummy_requests, which does not exist on main; they would fail at collection.
  • accuracy_core.py — the create_evaluator/prepared_evaluator refactor has zero callers even on the Rubin branch.
  • Both test-db YAMLs — their only additions reference test files that do not exist on main.

A new allowed_values member on a telemetry-captured categorical field is a manifest change, so telemetry/privacy CODEOWNER review applies. tests/unittest/api_stability/references/llm.yaml records only the annotation, not the Literal members, so no api_stability reference update is needed.

Known limitations

  • The CuTe DSL shim in tensorrt_llm/__init__.py runs at import time for every install. It is guarded by try: import cutlass.cute / except ImportError: return plus hasattr checks and should be inert when the pinned DSL already exposes the modern names, but that could not be confirmed here and it is on the hot path for all users.
  • dwdp/setup.py and modeling_kimi_k25.py are multi-GPU paths that could not be exercised.

PR Checklist

  • PR description explains what and why
  • Follows TRT-LLM coding guidelines to the best of our knowledge
  • API change present (MoeConfig.backend enum member) — api-compatible, additive only
  • Telemetry/privacy CODEOWNER review required for the manifest change

🤖 Generated with Claude Code

@reasonsolo
[None][chore] Core: MoE backend enum, deepgemm JIT warmup, postproc m…
50dc975 
…etrics, cleanup

Ports the non-Rubin-specific core changes from the internal Rubin branch.

Public API / config:
- Add "CUTEDSL_FC12" to MoeConfig.backend. The backend implementation
  (CuteDslFc12FusedMoE + moe_resolution registration) lands separately;
  this only opens the config gate. llm_args_golden_manifest.json is
  updated to match.

Autotune / JIT warmup:
- Add deep_gemm_jit_warmup_buckets() and use it for the two DeepGemm
  runners whose only purpose is to drive JIT warmup. A step-16 M grid over
  the whole range is required: the SM100 layout heuristic selects on
  ceil_div(m, block_m) with every candidate block_m a multiple of 16, and
  the last-wave-utilization tie-break keeps oscillating at high M, so a
  coarse high-M band silently skips layouts that then compile mid-inference
  (nvcc fork under the GIL, stalling every attention-DP rank).
  Fp8BlockScalingGemmRunner also gets exclude_from_cache so a warm disk
  cache cannot short-circuit the warmup.

Executor:
- Propagate decoding_iter, avg_decoded_tokens_per_iter and cached_tokens
  through PostprocWorker.Output, so a result served by a postproc worker
  reports the same metrics as the in-process path.
- When TRTLLM_WORKER_DISABLE_GC=1, also disable dynamo's post-compile
  gc.collect(1): with automatic GC off it walks every object allocated
  since the previous compile, costing seconds per recompile.

Models:
- Kimi K2.5: replicate the vision tower whenever cp_size > 1. Helix carries
  its parallelism in cp with tp_size=1, so the existing tp-only check never
  fired and the folded world size collapsed below the rank range.
- dwdp: resolve the MoE wrapper model-agnostically, so Kimi K3's
  block_sparse_moe/routed_experts spelling is handled alongside DeepSeek's
  mlp/experts.

Misc:
- CuTe DSL compatibility shim for legacy cute.core.ThrCopy/ThrMma and
  cute.make_fragment, needed by QuACK and Transformer Engine against the
  pinned CUTLASS DSL.
- Drop the unused <ATen/native/cuda/Resize.h> include from six thop GEMM
  translation units.
- Do not fail collection of tests/integration/defs when torch._inductor is
  unavailable.
- Unit test for submit.py's replace_env_in_file.

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
@reasonsolo
[None][fix] dispatch CuteDSL MoE backends by subclass in create_moe_b…
ec8d259 
…ackend

PR5 adds `CUTEDSL_FC12` to the `MoeConfig.backend` literal, but
`create_moe_backend` matched the CuteDSL branch by class identity:

    elif moe_cls in (CuteDslFusedMoE, CuteDslB12xFusedMoE):

`CuteDslFc12FusedMoE` subclasses `CuteDslFusedMoE` and takes the same
argument set, so selecting the newly-legal backend string would fall
through to the `Unsupported moe backend` raise.

Match by `issubclass`, as the DeepGEMM and MegaMoE branches below
already do for the same reason. Every class dispatched after this
branch derives from `MoEImplBase` / `MoE`, so the wider match steals
nothing; with no CuteDSL subclass on this branch the change is inert
until the FC12 backend lands.

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
@reasonsolo
[None][fix] add the fp8_prequantized_swap_ab_gemm op
88e1d89 
The pre-quantized FP8 path in Linear.apply selects
torch.ops.trtllm.fp8_prequantized_swap_ab_gemm whenever the activation
scale is int32 on SM100f with deep_gemm enabled, but no registration for
that op exists. Without it every DeepSeek V4 run on GB200/GB300 aborts
during executor init:

    AttributeError: '_OpNamespace' 'trtllm' object has no attribute
    'fp8_prequantized_swap_ab_gemm'

Unlike fp8_swap_ab_gemm, the activation and its packed scale both carry the
M dimension, so the tuning config ties input 1 dim 0 to input 0 dim 0, and
the runner restores the MN-major packed-scale stride that DeepGemm requires
after the autotuner recreates the constrained integer tensor contiguously.

Verified on GB300: the DeepSeek V4 DSpark disaggregated accuracy guard
(1p1d, DEP4) reaches GSM8K 96.21 with this op present and fails at executor
init without it.

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
@reasonsolo

reasonsolo commented Sep 16, 2026

Copy link
Copy Markdown
Collaborator Author

Added trtllm::fp8_prequantized_swap_ab_gemm (torch_custom_ops.py, +98).

Why this PR: Linear.apply's pre-quantized FP8 path selects this op whenever the activation scale is int32 on SM100f with deep_gemm enabled, but no registration existed. #19184 carries that caller (modules/linear.py:1255), so #19184 requires this PR — merging it alone reproduces:

AttributeError: '_OpNamespace' 'trtllm' object has no attribute 'fp8_prequantized_swap_ab_gemm'

Two details that differ from the sibling fp8_swap_ab_gemm: the activation and its packed scale both carry the M dimension, so the tuning config ties input 1 dim 0 to input 0 dim 0; and the runner restores the MN-major packed-scale stride DeepGemm requires after the autotuner recreates the constrained integer tensor contiguously.

Validated on GB300 (oci-jhb): the DeepSeek V4 DSpark disaggregated accuracy guard (1p1d DEP4, 2 nodes) reaches GSM8K 96.21 with this op present, and fails at executor init without it.

@reasonsolo reasonsolo mentioned this pull request Sep 16, 2026
Merged
3 tasks
@reasonsolo
[None][chore] drop CUTEDSL_FC12 now that the backend is not in this s…
104123b 
…eries

The CuteDslFc12FusedMoE backend has been removed from the MoE PR: the
custom op it drives was lost in the rebase, so it could never be
selected. Nothing in this series consumes the backend string, so the
enum value would be a `MoeConfig.backend` literal that resolves to an
empty candidate list.

Removes "CUTEDSL_FC12" from the literal and from the golden manifest,
and reverts ec8d259 -- the `issubclass` dispatch in
create_moe_backend existed only so the FC12 subclass would not fall
through to the "Unsupported moe backend" raise. With no CuteDSL
subclass in the tree the identity check is exact again.

Both come back with the backend and its op.

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
@reasonsolo reasonsolo mentioned this pull request Sep 22, 2026
Open
3 tasks
@reasonsolo

reasonsolo commented Sep 22, 2026

Copy link
Copy Markdown
Collaborator Author

Closing — the content of this PR is now consolidated into #19183, which also carries the sibling core/runtime work. Everything here that upstream had not already absorbed survives there: the disagg and KV-cache-v2 regression tests and the V2 draft-KV max_tokens derivation from this branch, and the DeepGemm JIT-warmup bucketing, executor postproc metrics and related core fixes from #19186. Branch is kept if anything needs to be recovered.

@reasonsolo reasonsolo mentioned this pull request Sep 22, 2026
Closed
5 tasks
@reasonsolo reasonsolo closed this Sep 22, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

@reasonsolo reasonsolo

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@reasonsolo