diff --git a/.github/application-test-inventory.json b/.github/application-test-inventory.json index 00807374c2..e6812f8846 100644 --- a/.github/application-test-inventory.json +++ b/.github/application-test-inventory.json @@ -9,6 +9,7 @@ "optional_collection_skips": { "tests/unit/collector/test_dsv4_megamoe_workload.py": "real torch required for tensor operations", "tests/unit/collector/test_helper_moe_distribution.py": "real torch required for tensor operations", - "tests/unit/collector/test_sglang_moe_ep_routing.py": "real torch required for tensor operations" + "tests/unit/collector/test_sglang_moe_ep_routing.py": "real torch required for tensor operations", + "tests/unit/collector/test_vllm_attn_kv_cache_layout.py": "real torch required for tensor operations" } } diff --git a/benchmarks/evidence/collector/op-probe-2026-10/README.md b/benchmarks/evidence/collector/op-probe-2026-10/README.md new file mode 100644 index 0000000000..bcde8e2252 --- /dev/null +++ b/benchmarks/evidence/collector/op-probe-2026-10/README.md @@ -0,0 +1,37 @@ + + +# Op-probe findings (2026-09 to 2026-10) + +This non-packaged bundle holds the serving observations that the collector +and generator changes for the vLLM 0.30.0, SGLang 0.5.21 and TensorRT-LLM +1.3.0rc29 pins cite as provenance. The op-probe harness (`opharness`) produced +them by booting the real serving framework per model and recording the +kernels and backends it selected. The harness lands separately. + +Citations in the code name either a finding key or a results file below. Grep +a finding key, such as `sm100_probe_failures_recheck_2026_10_04`, to reach its +statement. + +Every file here is an excerpt: `findings.yaml` keeps the 20 findings the code +cites, and each results file keeps `_meta` plus the models the citing comments +name. `_meta.summary` still counts the full matrix. The complete archive lands +with the op-probe harness. + +## Files + +- `findings.yaml`: curated conclusions, one entry per key. Each statement gives + its own versions, dates and hardware inline, and `applies_to` lists the + models it was observed on. Evidence paths inside statements refer to + workspace bundles that are not committed. +- `results/sm100/sglang-0.5.21.yaml`: B200 boot-and-run verdict per model on + SGLang 0.5.21 under the default generator render. +- `results/retests/sm100/sglang-0.5.21.moe_auto.yaml`: the B200 SGLang 0.5.21 + retest of the MoE models under an explicit `--moe-runner-backend auto`. +- `results/sm100/trtllm-1.3.0rc29.yaml`: B200 boot-and-run verdict per model on + TensorRT-LLM 1.3.0rc29. + +Findings describe framework releases and rot with them. Re-verify them on the +next version bump instead of extending them. diff --git a/benchmarks/evidence/collector/op-probe-2026-10/findings.yaml b/benchmarks/evidence/collector/op-probe-2026-10/findings.yaml new file mode 100644 index 0000000000..cacc8df43f --- /dev/null +++ b/benchmarks/evidence/collector/op-probe-2026-10/findings.yaml @@ -0,0 +1,913 @@ +# Excerpt (2026-10-30): only the findings that collector and generator code in +# this repository cite. The complete archive lands with the op-probe harness. +# OUTPUT: curated conclusions from probe campaigns (the facts archive's +# human layer). Each finding states its own evidence paths, versions and +# dates inline; applies_to lists where it was observed. Rots with +# framework releases — re-verify on version bumps. Inputs live in targets.yaml. +findings: + vllm_fp8kv: + statement: 'vllm 0.24/SM90: head-576 MLA has NO attention backend with fp8 KV; kv-auto(bf16) resolves DeepseekV2MLAAttention+SparseAttnIndexer + and runs. Door 1 engine kv dtype (fp8 ckpts) = config_gap, kv-auto A/B passes; door 2 ckpt hf_quant kv_cache_quant_algo:FP8 + (NVFP4 artifacts) = framework/platform gap, kv-auto still fails. Same two-door structure as trtllm. Evidence + facts/vllm_ab_kvauto_*.json, 2026-08-22. + 0.29.0 ADDENDUM (2026-09-24, GLM-5.3 onboarding): door 2 is now a SPELLING gap, not a missing kernel. With + kv auto the NVFP4 DSA checkpoints (GLM-5/5.1/5.2/5.3-NVFP4, DeepSeek-V3.2-NVFP4) resolve kv_cache_dtype to + the string fp8_e4m3 from the checkpoint kv scheme and the selector rejects every MLA backend + (FLASHMLA_SPARSE: "kv_cache_dtype not supported"); the explicit --kv-cache-dtype fp8 variant of the SAME + run loads FlashMLASparseImpl and generates (705890b31aa0, 4d34afefe419, ...). So on 0.29/SM90 these models + serve only with an explicit fp8 kv dtype — the golden vllm render passes none. Generator-relevant + (aiconfigurator should render --kv-cache-dtype fp8 for NVFP4 DSA on vllm 0.29) and upstream-worthy + (fp8_e4m3 vs fp8 in the sparse-MLA support check). Raw: archive/raw/67b8d5c6ce5e.json (Reasons dict). + GENERATOR FIX LANDED (owner "1，好的", commit ea0244a2): utils.vllm_dsa_kv_cache_dtype prescribes fp8 + keyed on architecture (GlmMoeDsaForCausalLM, DeepseekV32ForCausalLM) + the task''s nvfp4 GEMM quant mode; + module_bridge applies it to the worker payload (naive path has no quant fact). Left open for the SDK + owner: the perf task itself models these checkpoints with bf16 KV while serving runs fp8 KV either way.' + applies_to: + - zai-org/GLM-5-FP8 + - nvidia/GLM-5-NVFP4 + - nvidia/GLM-5.1-NVFP4 + - nvidia/DeepSeek-V3.2-NVFP4 + - zai-org/GLM-5.2-FP8 + - nvidia/GLM-5.2-NVFP4 + glm53_onboarding_2026_09_24: + statement: | + GLM-5.3 trio probed on all three backends after the rename (18 runs: + rendered + kv variant per backend; owner: "5.3 吧"). Matrix cells + (rendered config, framework mode): + vllm 0.29.0 GLM-5.3 (fp8) PASS FlashMLASparseImpl / Fp8MoE->triton_fused_moe; + GLM-5.3-BF16 PASS FlashMLASparseImpl / UnquantizedFusedMoE->triton_fused_moe; + GLM-5.3-NVFP4 FAIL (kv auto -> fp8_e4m3, no MLA backend; the + explicit fp8 variant PASSES with FlashMLASparseImpl — see the + vllm_fp8kv 0.29 addendum; generator should render + --kv-cache-dtype fp8 for NVFP4 DSA on vllm 0.29). + sglang 0.5.16 all three PASS (DeepseekSparseAttnBackend; NVFP4 MoE -> + ModelOptNvFp4FusedMoE->marlin on SM90); fp8_e4m3 variants + pass too. + trtllm rc23 GLM-5.3 PASS (DeepSeekFP8BlockScalesFusedMoE), GLM-5.3-BF16 + PASS; GLM-5.3-NVFP4 FAIL: the artifact is modelopt + MIXED_PRECISION (87 per-layer entries: FP8 dense mlp + projections + NVFP4 experts — GLM-5.2-NVFP4 was pure NVFP4) + and modeling_deepseekv3.py:1298 asserts "MIXED_PRECISION is + ambiguous" for every decoder layer -> unloadable on this + pin (real artifact property, not a dummy artifact; the + dummy keeps the top-level algo). fp8-KV variants of all + three fail at load like every DSA model on trtllm rc23 + (sparse FlashMLA asserts bf16 KV, see trtllm_kv_dtype_variants). + HARNESS GAP found: gen_dummy_models writes config.json only; the + tokenizer files every dummy dir needs were provisioned once by hand + (2026-09-19) and NO script does it for a new roster repo. GLM-5.3-BF16's + six runs first failed on "Couldn't instantiate the backend tokenizer" + (vllm/sglang at load, trtllm at generate: end_id needs a tokenizer); + fixed by copying tokenizer.json / tokenizer_config.json / + generation_config.json / chat_template.jinja from the GLM-5.3 sibling, + then re-probed. workflow_check already flags bare dirs ("variants + missing tokenizer"); onboard_model step 1 should own the copy. + MATRIX BUILDER DEFECT found and fixed the same hour: build_matrix took + ONE plan file per backend (the newest carrying the pin) and the version + from whichever backend's run came last in that file — the 18-run GLM-5.3 + plan produced an 18-cell matrix named *-1.3.0rc23.yaml for all three + backends. Now every plan file carrying (backend, pin) contributes, only + rendered (non-kv-variant) runs define cells, a passed cell is never + overwritten by a later crash/no-raw run, counts are recomputed, and the + cell's identity/kv_allocated comes from the rendered record (variant + records exist for the same repo x backend). Matrix: 98 models per + backend (95 + the trio); vllm 84 pass / 9 pass+custom / 5 fail (DSV4 + cells flipped fail -> pass+custom from the 6016-budget re-render), + sglang 86 / 6 / 6, trtllm 62 / 0 / 36. + applies_to: + - zai-org/GLM-5.3 + - zai-org/GLM-5.3-BF16 + - nvidia/GLM-5.3-NVFP4 + - instrument:probe_driver + - instrument:dummies + - framework:vllm-0.29.0 + - framework:sglang-0.5.16 + - framework:trtllm-1.3.0rc23 + - platform:sm90 + full_reprobe_framework_mode_2026_09_24: + statement: | + FULL RE-PROBE in framework mode (owner "现在跑吧"): every roster model on + the three pins, rendered config + kv-dtype variant — sglang 98+98, + trtllm 98+98, vllm 98+98 (4 Qwen3.8 runs skipped: the golden CLI + rejects the architecture). Old evidence retired under + facts/raw_superseded/ (187 eager vllm raws that already sat under the + current ids — the queue had skipped them as "done"; 291 pre-kv-tag raws + at the pins; 219 other-version records kept). records.jsonl now 713. + OUTCOMES (rendered / fp8-KV variant): sglang 93+91 ok — fails MiniMax- + M3-MXFP8 (mxfp8 unsupported), Kimi-K3 (+NVFP4, model module missing on + 0.5.16), Gemma-4-26B-A4B-NVFP4 (only gated SiLU), Llama-4-Scout (new + under graphs: sglang internal crash "Open an issue"), MiniMax-M2/M2.5 + fp8-KV only (bf16 output dtype assert); trtllm 62+51 ok — the known + SM90 wall (DSV4 pre-Blackwell assert, NVFP4 DSA fp8-KV, MiMo, M3 sparse + needs the kv cache type) plus every DSA fp8-KV variant; vllm 93+96 ok. + GATES on the new records: vllm 32 / sglang 16 / trtllm 13 aligned, + after three evidence corrections that framework mode forced: + - gemm_fp8 was paired with Qwen/Qwen3-32B-FP8-Static-PerTensor, which + vLLM 0.29 loads UNQUANTIZED (UnquantizedLinearMethod on every + linear; its config.json carries no quant_method and the modelopt + hf_quant_config.json in the dir is ignored) — never a valid fp8 + instance; the old "aligned" came from the fp8-KV variant record's + quant kernel. Re-paired with nvidia/Llama-3.1-70B-Instruct-FP8 + (ModelOptFp8LinearMethod, cutlass_3x_gemm_sm90_fp8) -> aligned. + - compute_scale (standalone scaled_fp8_quant kernels) has no serving + instance under torch.compile: the norm_quant / act_quant fusion + passes fold activation quant into Inductor kernels + (triton_*_fused__to_copy_clamp_cutlass_scaled_mm_*). Removed from + the gate; report kept next to the facts. The collector row is + eager-only truth — flag to the SDK owner (the op-level model bills + it separately). + - vllm MLA fp8 (DeepSeek-R1): under CUDA graphs the DeepGEMM + projection kernels sit under quant spans the MLA hint excluded -> + false collector-only deepgemm drift; hint widened (gemm|deepgemm|quant) + like the sglang DSA hint. + MATRIX (results/sm90, 100 models per backend): vllm 85 pass / 9 + pass+custom / 6 fail; sglang 87 / 6 / 7; trtllm 62 / 0 / 38. Flips vs + the eager matrix: GLM-5.1 / 5.2 / 5.3-NVFP4 vllm fail -> pass (the + generator now renders --kv-cache-dtype fp8; GLM-5-NVFP4 and + DeepSeek-V3.2-NVFP4 stay failed: bundled configs without + quantization_config, SDK data gap); Kimi-K3 (+NVFP4) vllm pass -> FAIL: + illegal memory access in kimi_k3 embed_input_ids during + compile_or_warm_up_model (CUDA-graph / compile warm-up) — --eager A/B + loads and serves (facts/full_reprobe/k3_eager.json), so it is a vLLM + 0.29 K3 compile/graph defect on SM90; kda_ctx keeps its verdict against + the eager A/B record (raw_file recorded). sglang MiniMax-M2 / M2.5 + fail -> pass under graphs; Llama-4-Scout pass -> fail (sglang crash). + Framework-mode probes: probe_eager False on all 588 new records. + applies_to: + - instrument:probe_driver + - instrument:path_diff + - framework:vllm-0.29.0 + - framework:sglang-0.5.16 + - framework:trtllm-1.3.0rc23 + - platform:sm90 + - moonshotai/Kimi-K3 + - Qwen/Qwen3-32B-FP8-Static-PerTensor + - nvidia/Llama-3.1-70B-Instruct-FP8 + followups_after_full_reprobe_2026_09_24: + statement: | + Owner directives after the full re-probe report: + (1) "统一所有的都对齐到64整数倍，别case by case了": the 64-alignment of + token budgets is now ONE rule in backend_config_mapping.yaml applied to + every backend (vllm max-num-batched-tokens, trtllm max_num_tokens, + sglang max-prefill-tokens and chunked-prefill-size); the vllm rule + formulas are plain again (commit 2446f6cc, mirrored into the aic + checkout the golden CLI renders with). Verified: sglang 5500 -> 5504, + trtllm 4500 -> 4544, vllm 6012 -> 6016. + (2) "关于hf config的quant，能不能看看其他字段": the KV fact is not in + config.json for the modelopt NVFP4 artifacts at all — the SDK loader + attaches the bundled hf_quant_config.json (kv_cache_quant_algo: FP8) + to the loaded config, so vllm_dsa_kv_cache_dtype now reads that merged + view. All five NVFP4 DSA artifacts (GLM-5 / 5.1 / 5.2 / 5.3, DeepSeek- + V3.2) render --kv-cache-dtype fp8; GLM-5-NVFP4 and DeepSeek-V3.2-NVFP4 + re-probed with the new render (see the matrix). + (3) "关于kimi k3，我感觉可能还有其他构造问题": the construction was + audited twice. Kimi-K3 names its layer kinds by 1-BASED index lists + (configuration_kimi_k3.is_kda_layer: `layer_idx + 1 in kda_layers`). + The depth cut kept the lists verbatim; by luck the 2-layer "rep" dummy + (original layers 1,2 -> 1-based 2,3, both KDA) still resolved to two + KDA layers, so the records' KDA evidence was real. My first remap + assumed 0-based lists and produced an all_kinds dummy whose two layers + both resolved to FULL attention — that variant "passed" the + framework-mode probe while running no KDA kernel at all (false pass, + caught by the record's kernel families). The generator now detects + the base from the config data (commits 85d1dd32 + 51158047): + all_kinds = one KDA (1) + one full-attention layer (2), no stale refs. + The crash itself: eleven A/Bs on the rendered config (max-num-seqs + 512) — every variant that runs a KDA layer fails (rep, 3-layer with + the dense layer, kda+kda, full+kda); CUDA graphs off passes; compile + off (graphs on) fails; piecewise-only graphs fail; text-only + (limit-mm image 0) fails; max-num-seqs 64 and 256 PASS. CUDA_LAUNCH_ + BLOCKING places the illegal memory access after graph capture, in + warmup_kernels -> _run_decode_step -> run_fullgraph — the decode + warm-up step at the generator's max_num_seqs=512 with CUDA graphs and + a KDA layer present. So: a real dummy defect existed and is fixed, and + independently vLLM 0.29's K3 (KDA) decode warm-up faults at batch 512 + on SM90 (256 passes) — the generator's 512 is what real deployments + get. Upstream draft 7 updated. The corrected all_kinds dummy (one KDA + + one full-attention layer) re-probed: the same warm-up crash — + attribution stands; kda_ctx keeps the eager A/B evidence (verdict + helper passes it explicitly). Side fact: the K3 fp8-KV variant is + rejected by vLLM at load — "Kimi-K3 fp8 KV cache requires an fp8 + prefill query; enable --attention-config use_prefill_q..." — an + extra flag the generator would have to render if fp8 KV were ever + chosen for K3 (it is not today). + (4) Dummy dirs: no step provisioned the tokenizer / custom-code files + (Kimi needs tiktoken.model + tokenization_kimi.py + modeling_*.py). + dummies.py now copies every auxiliary artifact file from + configs/aux_files/_/ (exported once from the existing dirs, + 100 repos), else an earlier variant of the repo, else the declared + targets.yaml dummy_overrides.tokenizer_from sibling; otherwise a loud + MISSING caveat. A fresh regeneration of all 196 variants from configs/ + is complete without hand copies (commit a2f34f3c). + (5) NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4 on trtllm rc23 ("no raw"): + the engine aborts right after model init — "CUDA runtime error in + ::cudaStreamDestroy(stream): an illegal memory access was encountered" + (tensorrt_llm TllmException, SIGABRT, so the probe cannot dump). + Same SM90 NVFP4 wall class as the other NVFP4 trtllm cells; log + facts/full_reprobe/nemotron_ultra_nvfp4_trt.full.log. + (6) Golden renderer switched from the aic checkout to THIS repository's + generator (venv_ais, `aiconfigurator cli generate` console script, + stamp = this repo's HEAD): vllm/trtllm renders byte-identical to the + aic ones for Llama-3.1-8B, sglang flags identical; Qwen3.8-2.4T-A95B + now renders (was "generator rejects" on aic). No more aic mirroring. + applies_to: + - instrument:dummies + - instrument:probe_vllm + - moonshotai/Kimi-K3 + - nvidia/Kimi-K3-NVFP4 + - framework:vllm-0.29.0 + - platform:sm90 + vllm_0300_upgrade_2026_09_27: + statement: | + Owner decision 2026-09-27: move the vllm pin from 0.29.0 straight to + 0.30.0 (vllm/vllm-openai:v0.30.0, the newest image; 0.29.x patches do + not exist), collector included — the 0.29 pin upgrade is abandoned and + its records left the PR (kept in the workspace, facts/local_only_results). + The upgrade_op workflow ran end to end as the stability test of the + harness; every step's evidence: + PROBES (224 vllm runs, framework mode, rendered + fp8-KV variants, + ~1.1 min median per run on the dummies, 50 min on six GPUs): 124-cell + final matrix 110 pass / 1 pass+custom / 13 fail (12 generator rejects, + 1 capacity: Qwen3.8-2.4T-A95B bf16 — no faithful cut fits one 139.8 GiB + GPU; its FP8 and NVFP4 artifacts pass at depth4). Two first-pass OOMs + on the fp8/NVFP4 2.4T cuts and on two fp8-KV variants were transient + GPU sharing with the concurrent capture wave and cleared on re-probe — + probes and captures must not share a GPU. Verdict changes vs 0.29: + DeepSeek-V4.1-Flash and its NVFP4 artifact now PASS (0.30 registers + DeepseekV41); Qwen3.8-2.4T-A95B-FP8 capacity -> pass (depth4 fits). IDENTITY: across the 94 repos + probed on both versions the identity-role (backend, role) sets are + IDENTICAL — no attention / gemm / moe / linear-attention / dsa / quant / + mhc kernel family changed between 0.29.0 and 0.30.0; 47 repos differ + only in sampler / scheduler glue kernels (0.30 drops + _convert_req_index_to_global_index / _compute_slot_mappings / + _get_num_sampled_and_rejected, adds cute-DSL cu_blocks conversion and a + symbol literally named "kernel"). New kernels labeled: DSV4 + _token_request_mapping_kernel (infra), Nemotron fused + _relu_squared_static_fp8_quant_kernel (0.30 fuses the activation into + the quant), DSV4.1 _hc_head_reduce_store_kernel (mhc). RESIDUE left on + purpose for the owner's granularity call: DSV4.1's engram module + (_engram_lookup_kernel, _hash_ids_kernel, _fused_engram_post_wkv_kernel) + and a DSV4 cute-DSL kernel_cutlass_device_kernel_tensorptrf32... whose + role is unattributed. The Qwen3.5 rows' attention column reads FA3 + + the GDN decode kernel instead of ChunkGatedDeltaRule + the same kernel: + 0.30 emits no wrapped attention span there, kernels unchanged. + COLLECTOR ON 0.30 (the harness's reason to exist): 21 captures re-taken + in the 0.30 image with versioned names (opcov_@0.30.0.json). The + first wave found the collector broken in two places — exactly the + metadata-parity class layer_permissions.md warns about: (a) + CommonAttentionMetadata lost the private _seq_lens_cpu / + _num_computed_tokens_cpu fields, the seq_lens_cpu property and the + num_computed_tokens_cpu property (10 MLA/DSV4 captures invalid); (b) + is_flashkda_supported() gained recurrent_state_dtype (kda_ctx invalid). + Fixed in collector/vllm/utils.py, collect_kda.py, collect_msa_module.py + with citations to the 0.30 population sites (v1/worker/gpu/attn_utils.py + :456-475, CommonAttentionMetadata.compute_num_computed_tokens, + models/kimi_k3/nvidia/kda.py:455-495 + :665, MambaStateDtypeCalculator. + kda_state_dtype -> float32 for the default mamba_ssm_cache_dtype); + framework_manifest vllm default -> 0.30.0 (digest recorded); the + __compat__ upper bound of all 11 vllm collector modules -> 0.30.0; + manifest tests updated. Re-captured: 19/19 declared gates ALIGNED + (compute_scale stays explained: the quant is still fused on 0.30, the + governed exception now scopes 0.30.0). Citation audit debt: only the + three broken sites were re-verified against 0.30 source; the other + @0.29.0 citations in the collector still hold by behaviour (their + captures align) but were not re-read. + CUSTOMIZATIONS (bare-default A/B at 0.30.0, results/retests/sm90/ + vllm-0.30.0.yaml): Nemotron-H-56B gpu-memory-utilization 0.6 still + needed (bare OOMs in the mamba profile run as on 0.29); the DSV4 + moe_backend=marlin prescription for the four fp4-experts checkpoints + is now DROPPABLE — 0.30's runner auto selects the same + marlin_moe_wna16 kernels the prescription forced (identity identical), + so all eight DSV4 vllm prescriptions were dropped from targets.yaml and + the eight repos re-probed under the bare render (their cells move from + pass+custom to pass). + KDA PIN RETIRED (owner 2026-09-28, "kda可以到0.30吗，不pin了"): the + manifest's kda family pin on the kimi-k3 preview image is gone; the + family resolves to the 0.30.0 default like every other vllm op + (collect_kda __compat__ 0.30.0). SMOKE ON 0.30 (collect.py --backend + vllm --smoke, one vllm/vllm-openai:v0.30.0 container, 4 random cases + per op, evidence facts/collector_smoke_0300/{main,kda0300,mini,full}): + the first passes exposed three more breaks, all fixed: (a) + is_fused_kda_decode_supported() also gained recurrent_state_dtype + (kda.py:665-680@v0.30.0; float32 as for the prefill probe) — kda decode + cases failed with TypeError; (b) finalize rejected the run — "staged + table scale_matrix_perf.txt has no selected producer": collect.py binds + every staged table to the registry entries declaring it and a producer + checkpoint owns exactly one table, while the compute_scale collectors + hard-coded a second table (PerfFile.SCALE_MATRIX registered 2026-09-21 + never fixed that). First fix (863bfc75) split the op into a + `scale_matrix` registry producer; OWNER DECISION 2026-09-28 ("做兼容 + 多表"): the executor supports MULTI-TABLE PRODUCERS instead — + OpEntry.extra_perf_filenames declares every further table one run_func + writes, the executor hands them in (run_computescale(..., + extra_perf_filenames)), provenance collections carry them, and every + finalize/transaction binding (_registered_checkpoint_tables, + producer plan, perf-transaction journal `tables`, sidecar table + ownership, collection_meta) treats a checkpoint as owning a SET of + tables. compute_scale is back to one op writing both tables from one + measurement; the split op and its YAML activations are gone. Verified: + mini smoke — one checkpoint vllm.compute_scale.json, computescale_perf + + scale_matrix_perf finalized with the same case_plan_hash, exit 0; + (c) the vllm image ships no pyarrow — the launchers pip-install it + before finalize. Full smoke (24 ops incl. the interim split op): 95/96 + cases pass (the one failure is the classified kernel-limit raise for + msa_generation_module at 6 indexer heads), 23 tables finalized as + parquet with a provenance sidecar, exit 0. + DATA: the collector default is 0.30.0 now and the collector finalizes + end to end on 0.30, but no 0.30 perf data campaign has run — the SDK + still ships vllm 0.24.0 rows; the campaign is the next (collector-side) + step. + applies_to: ['framework:vllm-0.30.0', 'instrument:probe_driver', 'instrument:path_diff', + 'instrument:workflow_check', 'collector:vllm', 'platform:sm90', deepseek-ai/DeepSeek-V4.1-Flash, + nvidia/Qwen3.8-2.4T-A95B-NVFP4, moonshotai/Kimi-K3, deepseek-ai/DeepSeek-V4-Flash, nvidia/Nemotron-H-56B-Base-8K] + rtx_systems_have_no_generator_facts_hardware_profile: + statement: 'GENERATOR GAP (found while re-probing V4.1 on sm120, 2026-10-01): model-default cli args from + generator/facts/models.yaml (e.g. V4.1 block-size 64, trust-remote-code, no-enable-flashinfer-autotune; V4 block-size 256) + are NEVER applied when --system is an RTX PRO system. request_resolution.hardware_key_for_system() has no entry for + rtx_pro_6000_server (map covers h100/h200/h20/b200/b300/gb200/gb300 only) and falls back to the raw name; resolve_facts() + raises KeyError("Unknown hardware profile rtx_pro_6000_server ... [b200, b60, gb200, h100, h20, h200]") and + pipeline.run_pipeline swallows it ("Fact resolution failed; continuing without facts", facts=None), so apply_facts is a + no-op. Every sm120 golden render in this campaign (archive/golden/*/render.log carries the KeyError) therefore ran + without model defaults; the harness has been compensating with --generator-set params.agg.trust_remote_code=true in + targets.yaml. The upstream "drop the harness override" for V4.1 block-size (a04c849c) is correct on H20 and silently + inert on RTX. Candidate fix (generator, owner approval): add an rtx_pro hardware profile to facts/hardware.yaml and map + rtx_pro_6000_server (and the future RTX PRO 5000 system) in _SYSTEM_TO_HW; or make unknown hardware degrade to + model-only facts instead of none.' + applies_to: + - platform:sm120 + - instrument:generator + - instrument:probe_driver + sm120_collector_gemm_fp8_static_row: + statement: 'COLLECTOR-SIDE DISPOSITION of the gemm_fp8 divergence (serving-parity audit + owner decision, 2026-09-30). + Audit: vLLM 0.30.0 selects the fp8 scaled-mm kernel by activation scheme — FlashInferFP8ScaledMMLinearKernel is + first in _POSSIBLE_FP8_KERNELS (kernels/linear/__init__.py:421-430), is_supported needs CC >= 100 + (scaled_mm/flashinfer.py:50-51), can_implement needs per-tensor activation AND weight scales (:56-63). ModelOpt + FP8 checkpoints (quant_algo FP8, static input_scale; modelopt.py:2333) satisfy it -> serving log "Selected + FlashInferFP8ScaledMMLinearKernel for QKVParallelLinear/RowParallelLinear/MergedColumnParallelLinear" + (facts/ab_sm120/da1900656836_baseline_log.log) -> cutlass MainloopSm120TmaWarpSpecializedBlockwiseScaling + + cuBLAS sm89_xmma at M=1. The collector''s `fp8` row is Fp8Config(activation_scheme="dynamic") -> can_implement + False -> CutlassFP8ScaledMMLinearKernel (op_smoke confirms). Reproduced in isolation: static vs dynamic Fp8Config + on one shape selects FlashInfer vs Cutlass (facts/ab_sm120/opcov_gemm_fp8_static_probe.json). fp8_static was NEVER + a collected row: the SDK derives it from the dynamic fp8 row + compute_scale + scale_matrix + (perf_database._gemm_key_names), which presumes one shared GEMM kernel — true below CC 100, false on SM100/103/120. + Every per-tensor FP8 checkpoint in the roster is ModelOpt static, so the gemm_fp8 gate never had a faithful + serving reference (sm90 aligned by coincidence). OWNER DECISION (tianhaox): do NOT add a collected fp8_static row + (a trial implementation was validated — op_smoke kernel_source=FlashInferFP8ScaledMMLinearKernel, gate aligned — + and then reverted); handle it as a governed approximation: collector/evidence_exceptions.yaml standards_exceptions + entry "fp8_static derived ... CC>=100" (scope gemm_perf / vllm 0.30.0 / CC>=100 systems, expires 2026-12-31), + gemm_fp8_Llama-3.1-70B-FP8 stays an explained deviation (facts/pathdiff/explained_0300). MLA fp8-KV on sm120 + (mla_gen_fp8 capture crashes like serving) is a framework-version gap with no YAML home (failure_handling.md): the + case fails classified and re-tests on the next pin bump; no registry unverified_sms mark since the bf16-KV rows + of the same op are validated.' + applies_to: + - platform:sm120 + - platform:sm100 + - platform:sm103 + - framework:vllm-0.30.0 + - instrument:collect_gemm + dsv41_dummy_cut_and_vllm_block_size_2026_10_01: + statement: 'CORRECTION of the DeepSeek-V4.1 rows in vllm_0300_upgrade_2026_09_27 and of the 09-30 support probe: the + V4.1 dummy those verdicts ran on was not V4.1. The generic depth cut kept original layers [0,1] — two sliding-window + layers (compress_ratio 0) — with kv_source_layer_ids / index_source_layer_ids / candidate_source_layer_id / + engram_layer_ids left verbatim (keys end in `_layer_ids`; the stale-ref post-check only knew `layers`), so no + compressor, indexer, compressed-KV attention or candidate path was ever exercised. FIX (owner request 2026-09-30, + "put the engram shrink in the dummy handling, no need for full size"): dummies.py `dsv41` adapter, routed by + targets.yaml dummy_overrides.family. Rep cut = original layers [0,1,2,3,20,21,24] (swa; engram swa; ratio-2 kv/index + source + consumer; candidate/ratio-1 source + consumer; candidate-masked indexer), rep_min = [1,2,3,20,21] as the + capacity fallback. Every pointer is renumbered (kv sources [2,4], index sources [2,4,6], candidate 4, engram [1]) and + the adapter raises if a compressed layer has no source at or below it (the frameworks assert exactly that: + vllm 0.30.0 models/deepseek_v41/attention.py:283-299, sglang main models/deepseek_v4.py:4316-4322). Engram hash + tables shrink by lowering engram_vocab_size 16000000 -> 100000 and recomputing engram_num_embeddings with the + frameworks'' prime-sum rule (vllm common/engram.py:186-199, sglang layers/engram.py:158-175; reproduces the + shipped 384006168/384016682 exactly): 2.4M rows = 0.6 GB per layer instead of 98 GB. Row count is memory only — + hash/gather/gate kernels and per-token work are unchanged; 100000 keeps buckets-per-head above the compressed + vocab (99092). The stale-ref check now covers `*_layer_id(s)` keys and checks nested sub-configs against their own + depth (vision_config, Inkling mtp_config). SIDE EFFECT: Inkling''s local_layer_ids (55 of 66) is a layer-kind + axis, so its generic variants became all_kinds [0,5] + single-kind with the list remapped (Inkling is a + generator reject on every backend, no record changed). RESULTS on the faithful cut (H20/sm90): vllm 0.30.0 + deepseek-ai/DeepSeek-V4.1-Flash FAILS at load on its golden render — the generator''s deepseek_v4 fact + (facts/models.yaml, matched through the `deepseek-ai/DeepSeek-V4*` path glob) renders block-size 256, and V4.1''s + compressed layers use 64-token kernel blocks, so the 256-token manager block must split 4-way, which the BLHNC + layout cannot do once a KV tensor is shared by several layers ("block stride 275904 != page 74752", + kv_cache_interface.py:373). Depth-independent: the real 40-layer deployment fails the same way. Rescue as the + framework''s own message suggests: `--generator-set params.agg.tokens_per_block=64` -> PASS (rendered and fp8-KV; + A/B in facts/../archive/superseded_dummies/2026-10-01_dsv41_inkling/raws_block256). The nvidia NVFP4 sibling is + NOT matched by that fact (arch DeepseekV41 / model_type deepseek_v41 / nvidia path), renders no block-size and + passes at vllm''s default — an inconsistency between two artifacts of one model that is the generator''s to + resolve. Matrix vllm-0.30.0: 110/1/13 -> 109/2/13 (V4.1 fp8 pass -> pass+custom). Identity now includes + DeepseekCompressor, DeepseekV4Indexer (deep_gemm sm90_fp8_mqa_logits / paged + sched metadata, + cp_gather_indexer_k_quant_cache, topKPerRowPrefill, cooperative_topk) and the candidate filter + (_block_scores / _candidate_flags / _mask_candidates / _store_candidates). RESIDUE for the owner''s granularity + call (unlabeled by the sm90 taxonomy on purpose): compressor glue (_fused_save_compress_norm, _rope_quant_insert, + _ring_slot_mapping), indexer-K store (_indexer_k_norm_rope_quant_store), the candidate kernels, engram + (_hash_ids / _engram_lookup / _fused_engram_post_wkv / _gather_lookback) and the cute-DSL LL bf16 router-gate + GEMM. sglang main (nightly-dev-20260930-cae69be5, exploratory, not a pin): all four V4.1 runs PASS (fp8 rendered, + fp8_e4m3 KV, NVFP4 x2 with the dsv4 marlin customization — runner=auto refuses NVFP4 on SM90 as for V4): backend + DeepseekV4AttnBackend, pool DeepSeekV4TokenToKVPool fp8, sparse_prefill_fwd + FlashMLA sparse decode (same sm90 + sparse kernels as vllm), but the indexer runs an fp4 index-K path (_quantize_fp4_indexer / _fp4_index_logits / + _store_fp4_index_k_cache) instead of vllm''s DeepGEMM fp8 mqa_logits, engram as _engram_hash/_gather/_gate. Before + the dummy fix the same run OOMed allocating the full engram table (183 GiB fp16 temp in initialize_dummy_weights). + sglang 0.5.20 (newest release) and trtllm 1.3.0rc29 do not support V4.1 (no model/config class), unchanged. + Probe-side compatibility for sglang main (probe_sglang.py): ServerArgs is a msgspec Struct, bench_one_batch moved + to sglang.benchmark.one_batch, config initializers take no args after publish(server_args); the generator still + renders `--cuda-graph-bs` which main splits into -decode/-prefill (one-off queue patch, argv stale vs plan by + design). COLLECTOR GAPS (unchanged): dsv4 case declarations for compress ratio 2/1 source/consumer layers and the + candidate filter; engram op; SDK composition by source/consumer layer counts.' + applies_to: + - deepseek-ai/DeepSeek-V4.1-Flash + - nvidia/DeepSeek-V4.1-Flash-NVFP4 + - baerquant/Inkling + - framework:vllm-0.30.0 + - framework:sglang-main-2026-09-30 + - instrument:dummies + - platform:sm90 + hopper_sglang_collector_port_0514_to_0521_2026_10_01: + statement: 'SGLANG COLLECTOR PORT 0.5.14 -> 0.5.21 + GATE BACKFILL (H20/sm90, 2026-10-01; owner decision 2026-10-01 ''跟到0.5.21''). + PIN: framework_manifest sglang default 0.5.14 -> 0.5.21 (lmsysorg/sglang:v0.5.21, index digest b1259f3e..., cu130 tag + same digest); the kda (kimi-k3 branch 0.5.16) and msa (v0.5.16) family pins are retired because 0.5.21 ships kimi_linear.py + / minimax_m3.py and both families smoke+gate clean on it; the three Qwen3.8 model pins stay at 0.5.17 (the 0.5.21 serving + matrix fails Qwen3.8-2.4T-A95B/-FP8 at boot). Every sglang module now declares __compat__ with 0.5.21 (validated set + = {0.5.14, 0.5.21} resp. {0.5.14, 0.5.17, 0.5.21}; intermediate releases excluded, unvalidated). COMPAT LAYER (collector/sglang/runtime_compat.py): + ServerArgs is a msgspec Struct (publish(server_args, role, ranks) then zero-arg initialize_moe/fp8/fp4 config; derived + args behind arg_groups resolving_view -> resolved_arg(); read-only after resolution -> get_context().override(source, + ...) for the CSA-context token cap), bootstrap.init_parallel_runtime + init_layer_runtime before ModelRunner, ensure_offline_runtime_published() + for mock backends, kv_index_translator / decode_num_tokens_per_req / kv_cache_dtype_str on the mock runners, import + relocations (jit_kernel -> kernels.ops, fp8_kernel -> kernels.ops.quantization, flash_mla -> sgl_kernel.flash_mla, communicator + -> layer_boundary.adapters, configurator._profile_available_bytes via model_runner.init_kv_cache_configurator()), native + DeepseekV4Config, FlashMLA sched-meta API detection. KERNEL-LANE FORKS (serving dispatch changed, all probed from the + framework source, not the version string): (1) GDN prefill on SM90/SM100 = FlashInfer chunk_gated_delta_rule (+gdn_prefill_qkv_prepare_fwd) + when gdn_backend.flashinfer_gdn_prefill_default would pick it (fp32 state on SM90, 128/128 head dims; Triton FLA otherwise) + and GDN decode conv = fused_qkvzba_causal_conv1d_update_contiguous (SGLANG_ENABLE_GDN_DECODE_FUSED_PROJ_CONV default + on) with the framework''s own eligibility fallback; gdn_ctx/gdn_gen gates went DIVERGED -> ALIGNED. (2) KDA decode for + lower-bound layers (Kimi-K3) = plain fused_sigmoid_gating_delta_rule_update (kda_backend.py:737-786), not the packed + kernel; kda_gen gate DIVERGED -> ALIGNED. (3) DSV4 paged-MQA schedule = sglang JIT get_paged_mqa_logits_metadata (paged_mqa_metadata_tiny_kernel, + SGLANG_OPT_USE_JIT_INDEXER_METADATA default on) instead of deep_gemm''s sched kernel; (4) DSV4 topk calib entry points + renamed topk_transform_512{,_v2} -> topk_transform_paged{,_v2} (same v1/v2 lanes). MEMORY PLANS: DSV4 HCA on 0.5.21 + prices a ratio-128 full token at ~64 B so the framework-derived pool on an empty H20 reaches ~2e9 tokens and the int64 + full_to_swa_index_mapping (not priced by DSV4PoolConfigurator) OOMs -> HCA now caps max_total_tokens to the planned + shapes (sglang clamps to min(profiled, cap)); MSA plan-time filter now prices every loaded layer (2) -> the three 33.6M-token + decode cells are dropped at plan time instead of failing the worker capacity check. Conv int32 guards in gdn/kda stay: + 0.5.21 casts offsets to int64 yet the >=2**31 cells still fault on H20. SMOKE: op_smoke 30 ops in the v0.5.21 image + -> 28 OK; gdn (106/112) and kda (91/112) fail only their own guard cells (classified raises). GATES sglang 0.5.21: 28/28 + declared gates ALIGNED (15 prior + 13 backfilled: dsv4 csa/hca ctx/gen, csa/hca attn, paged_mqa_logits, glm5 dsa_attn/mqa/topk, + kda ctx/gen, mla_bmm); encoder_attn parked as EXPLAINED (facts/pathdiff/explained_sgl_0521): the sglang plan probe is + text-only, its Qwen3-VL record has no vision-tower kernels (vllm proves this family from its multimodal profile_run) + -> the one remaining undeclared sglang family until the probe exercises the vision tower. GATES trtllm rc29 backfill + (decision 3): dsv4 csa/hca gen ALIGNED after two collector ports — DeepseekV4Attention(aux_stream_dict=...) and the + Hopper fp8_ds_mla pool (KvCacheConfig.dtype + extra_attrs[''kv_cache_dtype''], which also keeps the module off the fused + Q-FP8/KV-norm prologue serving never runs on SM90); dsv4 csa/hca ctx ALIGNED once the module learned the pool dtype + (without extra_attrs[''kv_cache_dtype''] it took the fused prologue and mla_rope_inplace rejected the strided raw latent) + -> trtllm rc29 16/17 gates aligned, 4 DSV4 gates backfilled. Known gaps carried: trtllm moe fp8_block gate still DIVERGED + on routing (RenormalizeMoeRoutingMethod vs DeepseekV3Gate); trtllm collector manifest default still 1.3.0rc20 (only + the harness targets pin moved to rc29); test_fullnode real_checkpoint tests fail on HEAD before this work (logger None + at collect.py:1279), unrelated. Evidence: facts/upgrade_rc29/sgl_smoke_r{1..9}, sgl_gates_0521.txt, trt_gates_rc29_backfill.txt, + launch_sgl_caps_backfill.sh.' + applies_to: + - framework:sglang-0.5.21 + - framework:trtllm-1.3.0rc29 + - platform:sm90 + - module:collector + - instrument:collect_gdn + - instrument:collect_kda + - instrument:collect_dsv4_attn + - instrument:deepseekv4_sparse_modules + - instrument:collect_msa_module + - instrument:collect_attn + - instrument:collect_mla_bmm + - instrument:runtime_compat + - instrument:trtllm.collect_dsv4_attn + - Qwen/Qwen3-VL-2B-Instruct + - moonshotai/Kimi-K3 + - sgl-project/DeepSeek-V4-Flash-FP8 + hopper_trtllm_collector_pin_rc20_to_rc29_2026_10_02: + statement: 'TRTLLM COLLECTOR PIN 1.3.0rc20 -> 1.3.0rc29 + MOE ROUTING PARITY (H20/sm90, 2026-10-02; owner asked whether + the two leftovers of the sglang port could be closed). PIN: framework_manifest trtllm default rc20 -> rc29 (digest 124da3f1...); + the rc23 moe / sparse_attention / msa family pins are retired, their history kept in the manifest comment (rc20 moe + data still flows through the declared reuse.yaml channel; the rc23 mhc evidence exception keeps its own scope). Every + trtllm module already declared __compat__ >= rc19/20/21. SMOKE: op_smoke over all 23 trtllm registry ops in the rc29 + image -> 21 OK; moe fails only on its first plan case by the collector''s own classified raise (fp8_block needs 128-aligned + shards: hidden 1024 / inter 2688 at tp16) and mamba2 only on its bs=64 x 32768 cells (CUDA OOM, memory) — both pre-existing + behaviours, not rc29 drift. GATES rc29: 17/17 ALIGNED. The moe_fp8block DeepSeek-V3 gate had been DIVERGED since the + rc29 audit because collect_moe routed every non-min-latency case through RenormalizeMoeRoutingMethod (customMoeRoutingKernel) + while serving builds DeepSeek-V3-family / GLM / Nemotron-H MoE from DeepseekV3Gate.routing_method (noaux_tc grouped + top-k with the correction bias -> deepseek_v3_topk_kernel; modeling_deepseekv3.py:979-990, modeling_glm.py:273-285, + modeling_nemotron_h.py:236). The collector now reads the model row''s declared routing (cases/models/*_cases.yaml sglang_moe_routing_method_type=DeepSeekV3 + with num_expert_group / topk_group / routed_scaling_factor — framework-neutral checkpoint facts despite the prefix) + and builds the same gate routing with the resolved moe_backend; softmax-top-k models keep Renormalize. The router GEMM + (dsv3_router_gemm_op) stays outside the op as before. OPEN on trtllm: 7 undeclared gate families (encoder_attn, gdn_gen, + mamba2, mla_bmm, msa ctx/gen, ...) — gate_declared todo, same backfill mechanics as the sglang round. TEST HYGIENE: + collector/collect.py''s module logger defaulted to None until main() ran setup_logging, so in-process callers of the + checkpoint/provenance paths (fullnode finalization, unit tests run in isolation) crashed on logger.info; it now defaults + to logging.getLogger(''collector'') with a configured flag — mechanism-adjacent plumbing only, no behaviour change for + the CLI (main still installs the configured logger once). test_fullnode 8/8 standalone; the 30 test_collect_provenance_writer + failures are environment-only (no torch in this venv) and identical on the pre-change HEAD.' + applies_to: + - framework:trtllm-1.3.0rc29 + - platform:sm90 + - module:collector + - instrument:trtllm.collect_moe + - instrument:collect.py + - deepseek-ai/DeepSeek-V3 + - zai-org/GLM-5 + - sgl-project/DeepSeek-V4-Flash-FP8 + sm100_moe_backend_fact_rejected_2026_10_04: + statement: 'GENERATOR FACT BUG (B200 / sm100, 2026-10-04, first sglang 0.5.21 + trtllm 1.3.0rc29 identity probes on Blackwell; + Hopper never saw it): generator/facts/hardware.yaml b200 and gb200 profiles carry `moe_backend: {trtllm: WIDEEP, sglang: + deepep_moe}` and facts/apply.py apply_moe_backend fills it on EVERY MoE deployment (fill-if-absent, keyed on hardware only + — no EP size, node count or framework version). Neither value exists at the pinned versions: trtllm rc29 MoeConfig.backend + is Literal[AUTO, CUTLASS, CUTEDSL, CUTEDSL_FC12, TRTLLM, DEEPGEMM, DENSEGEMM, VANILLA, TRITON, MARLIN, MEGAMOE_*] (WIDEEP is a + 1.0-era literal; the rc29 template comment still lists it) -> pydantic literal_error before boot; sglang 0.5.21 + --moe-runner-backend rejects deepep_moe (argparse invalid choice). 156 trtllm + 150 sglang MoE renders (every MoE cell x kv + variant) died before any kernel ran; the as-rendered matrices read 29/0/95 (trtllm) and 32/3/89 (sglang). The h100/h200/h20 + profiles say `trtllm: CUTLASS` and leave sglang unset, which rc29/0.5.21 accept — that is why the sm90 campaign passed and the + fact was never exercised. The fact''s premise ("wrong MoE backend on Blackwell is a startup crash") is not what the A/B shows + (facts/ab_trt_moe_backend/, 12 probes on an empty B200): with moe_config dropped (rc29 AUTO) DeepSeek-R1, Qwen3-30B-A3B, + gpt-oss-120b and Qwen3-235B-A22B-NVFP4 all boot and run MoE kernels; explicit CUTLASS fails DeepSeek-R1 (CutlassFusedMoE + FP8_BLOCK_SCALES only supports SM90/SM120, got SM100); DEEPGEMM passes. Wide-EP in current TRT-LLM is an EP-size / + load_balancer / attention-dp deployment choice layered on CUTLASS|TRTLLM|DEEPGEMM, not a per-layer MoE backend literal. + FIX (local commit 8b96aa45b, owner integrates): the two Blackwell values are commented out with this rationale; profiles + keep `moe_backend:` so apply_moe_backend sees None. IDENTITY UNDER THE FRAMEWORKS'' OWN CHOICE: the same 306 cells were + re-probed with the fact removed from the rendered config (B200 workspace archive/retests/moe_auto/{raw,*.engine.yaml}; + table results/retests/sm100/.moe_auto.yaml): sglang 72 pass / 3 fail (Kimi-K3 bf16 `multi_ctas_kv_counter_buffer is + only supported when a trtllm-gen runner is selected`; Kimi-K3-NVFP4 and Qwen3.8-2.4T-NVFP4 `FusedMoE has no attribute + g1_scale_c`), runners FLASHINFER_TRTLLM 47 / AUTO 11 / FLASHINFER_MXFP4 7 / FLASHINFER_TRTLLM_ROUTED 3 / MARLIN 2 / TRITON 1 + / DEEP_GEMM 1, kernel families trtllm_gen_moe 57 / triton_fused_moe 12 / marlin 2; trtllm 61 pass / 17 fail (8 DSV4 — see + sm100_trtllm_dsv4_facts_unservable_2026_10_04; MiniMax-M3 x3 and Kimi-K2.5 quant parser gaps; MiMo-V2-Flash arch; + V4.1-Flash x2 transformers too old in the rc29 image; Kimi-K3 x2 no `fla` module), modules UnquantizedFusedMoE 23 / + NVFP4CutlassFusedMoE 20 / DeepSeekFP8BlockScalesFusedMoE 13 / FP8QDQ 2 / W4A8MXFP4MXFP8TRTLLMGen 2 / W4A16NVFP4Cutlass 1. + For sglang the fixed-fact render equals the override command byte-for-byte (150/150), so those raws ARE the cells'' evidence + (installed with the v2 fingerprint sidecars); for trtllm the fixed-fact render is NOT the override — see the next finding. + HARNESS DEFECTS SURFACED: (1) probe_sglang.py dies in argparse and writes no raw, so 150 cells graded "no raw (crashed before + dump)" with no error text (probe_trtllm records the pydantic error); the probe should dump an error record on argparse + failure. (2) _fail_cause had no rule for a rendered value the framework refuses — the trtllm cells read "framework gap"; + rule added: "generator fact rejected by framework" (literal_error | invalid choice). (3) render_golden''s cache key is + `git rev-parse HEAD`: an UNCOMMITTED generator edit re-renders nothing (the first fixed-fact emit silently served the WIDEEP + goldens) — key on a dirty-tree hash or on the facts files'' content. (4) tp4 jobs are emitted into the first GPU''s queue + (comment says run them alone) and were twice run from a stale script here; sglang''s tp memory-balance check also aborts a + tp4 launch when ANY of its GPUs is busy (Qwen3.8-2.4T-FP8: "memory capacity is unbalanced") — tp>1 cells need a + whole-group reservation in the queue plan. Re-run on umbriel-b200-021 with the group idle: sglang Qwen3.8-2.4T-A95B-FP8 tp4 PASSES (both kv variants; the results/sm100 matrix cell ''framework gap'' is that busy-GPU artifact), sglang Qwen3.8-2.4T-A95B bf16 tp4 flaked once (RemoteDisconnected) and passed on retry, trtllm Qwen3.8-2.4T-A95B-FP8 tp4 fails in the executor worker (CUTLASS fp8-block class; the matrix cell ''no raw'' is this).' + applies_to: + - platform:sm100 + - framework:trtllm-1.3.0rc29 + - framework:sglang-0.5.21 + - instrument:generator.facts.hardware + - instrument:generator.facts.apply + - instrument:probe_driver + - instrument:probes.probe_sglang + collector_impact: + - {op: moe, action: guard, ref: '95085922 collect_moe._raise_if_unverified_moe_lane: int4_wo opened on SM100/103 for 0.5.21 — the moe_auto retest here (Kimi-K2.5 FLASHINFER_TRTLLM / trtllm_gen_moe) is the evidence; lane_evidence.yaml now grades it'} + - {op: none, action: pipeline, ref: 'generator hardware fact (moe_backend deepep_moe) is the generator plane; collector cases unaffected'} + sm100_trtllm_dsv4_facts_unservable_2026_10_04: + statement: 'GENERATOR DSV4 FACTS ARE A HOPPER RECIPE (B200, trtllm 1.3.0rc29, 2026-10-04): the 8 DeepSeek-V4 trtllm cells + (Flash/Pro x {bf16, FP8, NVFP4, 0731/0813}) render `kv_cache_config: {dtype: fp8_ds_mla, tokens_per_block: 128}`. On SM100 + rc29 first rejects the page ("DeepSeek-V4 fp8_ds_mla KV cache requires tokens_per_block=256 on SM100, got 128"; cause + "generator DSV4 kv facts unservable on this SM"); with page 256 (archive/retests/tpb256{,_auto}/) the fp8_ds_mla SPARSE + GENERATION step has no FMHA library at all — TLLM_LOG_LEVEL=DEBUG (facts/dsv4_fmha_debug.log): CuTe DSL MLA "sparse + attention (DeepSeekV4Params) is not supported", FlashInfer trtllm-gen "sparse attention", Triton custom-mask "only context", + PrimTS "block-sparse inputs required" -> "No TRT-LLM attention FMHA library supports this request" (attention/backends/ + trtllm.py:2323). The servable sm100 configuration is kv dtype auto (bf16, DeepseekV4CacheManager) + tokens_per_block 256 + + MoE AUTO (facts/ab_dsv4_kvauto/, archive/retests/dsv4_sm100/): DeepSeek-V4-Flash-FP8 runs fmhaSm100f + QkvBfloat16 H512 PagedKvDense DynamicTokenSparse cubins, applyMLARopeAndAssignQKV, topKPerRowDecode, trtllm-gen routing. + GATES: the four trtllm dsv4 gates graded against that raw (--serving-raw) are DIVERGED by KV dtype — the collector''s DSV4 + lane follows the generator fact (E4m3 cubins, dsa_nsa indexer) while the only configuration rc29 can serve on sm100 is bf16 + — so on sm100 either the generator DSV4 facts move to {auto, 256} and the collector grows a bf16 DSV4 lane, or rc29 ships the + fp8_ds_mla sparse-gen kernel; until then DSV4/trtllm on sm100 is a platform gap, not a collector verdict. sglang''s DSV4 + cells and all 7 sglang dsv4 gates are ALIGNED on sm100 (0.5.21 defaults DSV4 kv to fp8_e4m3 itself).' + applies_to: + - platform:sm100 + - framework:trtllm-1.3.0rc29 + - instrument:generator.facts.models + - instrument:trtllm.dsv4 + - instrument:path_diff + sm100_sglang_trtllm_gate_verdicts_2026_10_04: + statement: 'COLLECTOR VERDICT, sglang 0.5.21 + trtllm 1.3.0rc29 on sm100 (B200, 2026-10-04; captures facts/pathdiff/ + opcov_@.json, verdicts results/pathdiff/sm100/-/). sglang: 14 ALIGNED (dsa fp8 ctx s4096 / gen, dsv4 + csa|hca ctx|gen|attn + paged_mqa_logits, glm5 mqa_logits + topk, gdn_ctx, mla_bmm_gen), 3 DIVERGED, 11 MISSING CAPTURE, 2 no + serving record (kda: Kimi-K3 does not boot on 0.5.21/sm100). trtllm: 7 ALIGNED (gdn ctx/gen, mamba2 ctx/gen, mla_ctx_fp8, + mla_gen_fp8, mla_bmm_gen — the three DeepSeek-V3 ones against the AUTO override raw, the render itself cannot serve + DeepSeek-V3 on sm100), 7 DIVERGED (attn fp8 ctx/gen, dsv4 x4, encoder), 8 MISSING CAPTURE. All 32 captures with a checked-in + recipe were taken (18 + 14); the trt_kv recipes hardcode the H20 layout `/work/ais/...` and needed the checkout mounted + there. BLACKWELL-SPECIFIC DIVERGENCES (all ALIGNED on sm90): (a) trtllm attn_ctx_fp8 / attn_gen_fp8 — collector/trtllm/ + collect_attn.py:171 hardcodes `tokens_per_block = 64` while the generator renders 32 for every trtllm deployment + (kv_cache_config.tokens_per_block: 32 in 208/224 sm100 renders); the sm100 cubin names carry the page (fmhaSm100a…P64 vs + …P32) so the mismatch is visible here and invisible in sm90 names — same class as the vllm block_size finding: the + collector should take the page from the deployment facts, not a constant. FIXED + VALIDATED (umbriel-b200-021, 2026-10-04): collect_attn now asks tensorrt_llm.llmapi.KvCacheConfig() for its default page (32; AIS_TRTLLM_TOKENS_PER_BLOCK is the A/B hook); re-captured in the rc29 image the collector runs fmhaSm100aKernel_QkvE4m3OBfloat16H128PagedKvCausalP32VarSeqQ256Kv128PersistentContext (ctx) and fmhaSm100fKernel_QkvE4m3OBfloat16H128PagedKvCausalP32MultiCtasKvCgaVarSeqQ8Kv256 (gen) = the serving cubins; the AIS_TRTLLM_TOKENS_PER_BLOCK=64 A/B reproduces the P64 divergence (facts/pathdiff/opcov_trt_attn_{ctx,gen}_fp8kv@1.3.0rc29{,.p64}.json in the 021 workspace). Re-graded on umbriel-b200-021 against fresh Llama-3.1-8B fp8-KV serving records (scoped re-probe, plan_*_scoped.json): attn_ctx_fp8 and attn_gen_fp8 are now ALIGNED (results/pathdiff/sm100/trtllm-1.3.0rc29/attn_*_fp8_Llama-3.1-8B.json), so trtllm reads 9 ALIGNED / 5 DIVERGED (dsv4 x4 by KV dtype, encoder). (b) sglang dsa_gen_bf16 — 0.5.21 + arg_groups/overrides.py resolves DSA kv auto to fp8_e4m3 when SM major >= 10 (bfloat16 on Hopper), so serving runs + QkvE4m3 cubins and the bf16 DSA lane is a non-case under the default deployment on Blackwell (dsa fp8 gates ALIGNED); + grade bf16 DSA gates on sm100 only against an explicit --kv-cache-dtype bf16 record — DONE on umbriel-b200-021 (archive/retests/dsa_bf16/, verdicts results/pathdiff/sm100/sglang-0.5.21/*.kvbf16.json): with --kv-cache-dtype bfloat16 both dsa_gen_bf16_DeepSeek-V3.2 and glm5_dsa_attn_GLM-5 are ALIGNED, so the collector''s bf16 lanes are right and the default-deployment divergence is purely the Blackwell kv-auto->fp8 rule. (c) sglang glm5_dsa_attn — + collector/sglang/glm5_dsa_sparse_modules.py runs the dsa_attn kernel as bf16 FlashMLA sm100 sparse_attn_fwd while serving + (kv auto -> fp8) runs TRT-LLM FMHA token-sparse cubins (fmhaSm100f QkvE4m3 HQk576HV512 …TokenSparse); the glm5 module + needs the fp8-KV lane the dsv32 module already has (the bf16 lane itself is ALIGNED under explicit bf16 KV). (d) encoder_attn both frameworks — op-hint issue as on vllm (owner + decision pending). OTHER REAL sm100 FAILURES in the matrices: MiMo-7B / MiMo-V2-Flash arch not registered (rc29), Muse- + Glimmer x2 and DeepSeek-V4.1-Flash x2 transformers too old in the rc29 image, MiniMax-M3 x3 sparse cache None, Kimi-K2.5 + quant parser, Kimi-K3 x2 no `fla` (trtllm); Kimi-K3 bf16 multi_ctas_kv_counter_buffer, Kimi-K3-NVFP4 / Qwen3.8-2.4T-NVFP4 + g1_scale_c (sglang). The 12 generator-rejects are identical to sm90 (roster repos aiconfigurator does not know).' + applies_to: + - platform:sm100 + - framework:sglang-0.5.21 + - framework:trtllm-1.3.0rc29 + - instrument:trtllm.collect_attn + - instrument:sglang.glm5_dsa_sparse_modules + - instrument:captures.trt_kv + - instrument:path_diff + collector_impact: + - {op: none, action: none, ref: 'verdict plane only (gates/captures); collector grids unchanged'} + sm100_probe_failures_recheck_2026_10_04: + statement: 'RE-CHECK OF EVERY sm100 PROBE FAILURE / NON-ALIGNED GATE FOR A CONFIG CAUSE (owner ask; umbriel-b200-037, fresh + workspace, Hub configs re-fetched). OURS (generator / harness / dummy): (a) 12 generator rejects per framework — roster repos + aiconfigurator does not know (identical on sm90); (b) trtllm 19 fp8-block MoE cells — rc29 template MoE default CUTLASS + (sm100_trtllm_moe_default_cutlass_2026_10_04); (c) trtllm DSV4 x8 — fp8_ds_mla/page facts (sm100_trtllm_dsv4_facts_unservable_2026_10_04); + (d) trtllm Qwen3.8-2.4T-FP8 tp4 — same CUTLASS class (sm90 pass); (e) encoder_attn DIVERGED on sglang AND trtllm — the serving + record carries 0 ops for Qwen3-VL-2B because the identity probe is text-only and the vision sidecar (probes/vision_.py) + exists for vllm only: harness gap, not a kernel finding; (f) **MiniMax-M3 x3 on trtllm (also failed on sm90, graded "quant parser + gap")** — DUMMY CUT BUG: 1.3.0rc29''s MSA cache manager does not read sparse_attention_freq; without explicit sparse_layer_ids it + allocates INDEX_KEY buffers for range(3, num_layers) (attention/backends/sparse/minimax_m3/cache_manager.py:303) while the model + builds sparse layers from the freq list (modeling_minimaxm3.py:363); our 2-layer "moe_sparse_attn" cut made both layers sparse at + positions the manager treats as dense -> idx_cache None -> "''NoneType'' object has no attribute ''shape''" in write_layer_caches. + The real 60-layer model has exactly 3 dense layers first, so serving never sees it; sglang reads the freq list and passed with the + same dummy. dummies.py variants_m3 now cuts the first three dense layers + the first sparse layer (4 layers): all six MiniMax-M3 trtllm cells (M3 / M3-MXFP8 / M3-NVFP4 x bf16 + fp8 KV) PASS on B200 with the MSA kernels visible + (SparseAttentionForwardSm100, _gqa_sparse_decode_kernel, fmhaSm100f cubins) — the three results/sm100/trtllm-1.3.0rc29.yaml + cells graded ''quant parser gap'' become pass once the matrix is rebuilt with the new dummy (same on sm90). FRAMEWORK (same on sm90 unless noted): (g) Kimi-K2.5 + trtllm — compressed-tensors W4A16 config (input_activations: None, mirrors the Hub config) dereferenced unguarded + (inputs_quant_config["strategy"]); (h) MiMo-7B / MiMo-V2-Flash arch not registered; Muse-Glimmer x2 + DeepSeek-V4.1-Flash x2 need a + newer transformers than the rc29 image; Kimi-K3 x2 trtllm image lacks `fla`; (i) **Kimi-K3 sglang bf16 (sm100 ONLY; sm90 pass)**: + Kimi-K3 has 96 attention heads; flashinfer''s trtllm-gen MLA decode refuses 64 < num_heads_q < 128 (mla/_core.py:3948), the runner + falls to cute-dsl, but sglang 0.5.21''s trtllm_mla backend passes multi_ctas_kv_counter_buffer unconditionally on Blackwell -> + "multi_ctas_kv_counter_buffer is only supported when a trtllm-gen runner is selected" (_core.py:4278). tp1 only: at tp>=2 the + per-rank head count is <=48 and supported — the probe''s single-GPU tp1 geometry is half the story; probe 96-head MLA models at tp2 + on sm100. (j) Kimi-K3-NVFP4 / Qwen3.8-2.4T-NVFP4 sglang "FusedMoE has no attribute g1_scale_c": modelopt_quant''s NVFP4 MoE weight + prep gates the flashinfer-trtllm branch on the GLOBAL get_moe_runner_backend() (modelopt_quant.py:2837) while apply() uses the + layer-cached resolved backend (:3096-3129) — under AUTO the two disagree and g1_scale_c is never created. A/B with an explicit + --moe-runner-backend flashinfer_trtllm clears it (facts/ab_k3nvfp4_runner: the cell then proceeds to the 96-head MLA error above), + so that flag is a valid customization once tp2 is used. BOOTSTRAP: the sequential roster fetch hit Hub HTTPErrors on 30 repos and + fell back to bundled copies (MiniMax-M3 flat vs VL layout; a sibling hf_quant without exclude_modules) — dummies.py now tolerates + both; fetch_inputs should retry before falling back.' + applies_to: + - platform:sm100 + - framework:sglang-0.5.21 + - framework:trtllm-1.3.0rc29 + - instrument:dummies + - instrument:fetch_inputs + - instrument:probes.vision + collector_impact: + - {op: mla_generation, action: guard, ref: '95085922 collect_mla.run_mla: FIXME(kernel-limit) classified raise for trtllm_mla decode with 64 < local heads < 128 on SM100/103 (item (i)); job 469988017 failed exactly the 362 num_heads=96 cases before this landed'} + - {op: none, action: none, ref: 'item (j) g1_scale_c: collector passes the runner explicitly, not reachable from the grid; dummy/fetch fixes are the probe plane'} + sm89_vllm_0300_probe_2026_10_04: + statement: 'FIRST sm89 (L40, 8 x 44.4 GiB usable, CC 8.9, driver 595.58) identity campaign, vllm 0.30.0 (image digest + 8a69ffad015f), generator = this branch, system l40s (needs the new l40s hardware profile: without it resolve_facts raised + ''Unknown hardware profile l40s'' and the pipeline swallowed it, the sm120 failure class). Matrix results/sm89/vllm-0.30.0.yaml: + 67 pass / 4 pass+custom / 53 fail (sm90: 110/1/13, sm120: 98/7/19). Fails, all caused: 17 capacity (44 GiB: Qwen3.5-122B/397B + depth-4 cuts, Nemotron-3-Ultra, Llama-4 Scout/Maverick - the Llama-4 failure was a dummy-staging gap (no preprocessor_config.json, + gated repo) fixed by the mirror files and is now capacity), 14 PLATFORM GAP no sparse-MLA backend: every DSA checkpoint + (DeepSeek-V3.2, GLM-5/5.1/5.2/5.3 incl. NVFP4) dies at attention-backend selection with use_sparse=True for kv auto + AND fp8 (the catch-all used to file these as the config cause ''ckpt-forced fp8-KV'': the exception line was truncated + before use_sparse), 12 generator rejects (same SDK gaps as sm90/sm120), 7 DeepSeek-V4: DeepGEMM hyperconnection.hpp:56 + ''Unsupported architecture'' at model init, 1 Kimi-K3-NVFP4 (_KimiGDNMergedColumnParallelLinear has no orig_dtype, same + on sm120). The 4 NVFP4 MLA checkpoints (DeepSeek-V3.1 / Kimi-K2.5 / K2.6 / K2.7-Code NVFP4) fail the SAME way as sm120 + - vLLM resolves modelopt NVFP4 to fp8 KV and TRITON_MLA fp8 decode needs 102400 B smem vs the 101376 B opt-in limit + that Ada shares with consumer Blackwell (Required 102400 / Hardware limit 101376); A/B: --kv-cache-dtype bfloat16 boots + (raws 6da3d656dea1, ee9f7f075eb4, dc45ccaead4c, 3924e9cb23ea), the fp8-KV variants stay red (platform floor). The sm120 + customization was extended to sms [sm120, sm89] on that evidence (results/retests/sm89/vllm-0.30.0.yaml). Identity on + Ada: dense attention FlashAttention-2 / FlashInfer (no FA3), MLA TRITON_MLA, MoE marlin/triton_fused_moe, bf16 GEMM + cuBLAS ampere_* kernels. Gates: 13 real, ALL ALIGNED (gemm bf16 + m1, fp8, moe bf16/mxfp4, mla ctx/gen/bmm, gdn ctx/gen, + attn ctx/gen, encoder, msa ctx/gen); 18 declared as sm89 platform floors with the cause in the verdict file; compute_scale + stays the explained governed approximation as on sm90. workflow_check upgrade_op vllm 0.30.0 sm89: 6/7 steps green, + failures_root_caused needs this entry.' + applies_to: + - platform:sm89 + - framework:vllm-0.30.0 + - nvidia/DeepSeek-V3.1-NVFP4 + - nvidia/Kimi-K2.5-NVFP4 + - nvidia/Kimi-K2.6-NVFP4 + - nvidia/Kimi-K2.7-Code-NVFP4 + - deepseek-ai/DeepSeek-V3.2 + - zai-org/GLM-5 + - nvidia/Kimi-K3-NVFP4 + - Qwen/Qwen3.5-122B-A10B + - Qwen/Qwen3.5-397B-A17B + - nvidia/Qwen3.5-122B-A10B-NVFP4 + - nvidia/Qwen3.5-397B-A17B-NVFP4 + - nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 + - nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-FP8 + - meta-llama/Llama-4-Scout-17B-16E-Instruct + - meta-llama/Llama-4-Maverick-17B-128E-Instruct + - instrument:probe_driver + sm89_sglang_0521_probe_2026_10_04: + statement: 'sm89 (L40) identity campaign, sglang 0.5.21 (digest b1259f3ea3). results/sm89/sglang-0.5.21.yaml: 65 pass + / 0 custom / 59 fail (sm90: 101/4/19). Fails: 14+7 capacity (7 = ''Loaded weights leave no GPU memory for the KV cache'' + at the planned mem-fraction: MiniMax-M3, GLM-5, Qwen3.5-122B, nvidia/Qwen3.5-397B-A17B-NVFP4 ...), 12 generator rejects, + 14+3 platform floors (sparse attention ''Sparse Attention Forward Kernel is only supported on SM90/SM100''; DeepSeek-V4 + sparse decode: DeepGEMM attention.hpp:496 asserts arch_major in {9,10,12} at CUDA-graph capture; MXFP dense GEMM / FP4 + quantization have no SM89 backend: flashinfer ''Invalid backend: 89''), 5 FRAMEWORK DEFECT: flashinfer 0.6.18''s cutlass + fused-MoE JIT does not compile for sm_89 (''tensorrt_llm::kernels::cutlass_kernels has no member Fp4Type'', ninja build + of fused_moe_89) - the Nemotron-3 family (Nano/Super/Ultra/Lightning), reproduced from the raw; 2 Qwen3.8-2.4T cells + died in sglang''s tp memory-balance check because the queue puts tp>1 cells on ONE GPU''s queue without reserving the + group (harness defect, rerun alone). Gates: 5 ALIGNED of the 6 real ones (attn ctx/gen, attn_ctx_hd512_gemma-4, gdn + ctx/gen), 23 declared sm89 floors (sparse, capacity for the V3-bound MLA/MSA/KDA/GLM-5 gates, DeepGEMM arch). encoder_attn_Qwen3-VL-2B + was DIVERGED for lack of evidence (the text-only probe never ran the vision tower, so the profile_run phase the gate + asks for was empty); CLOSED in the review pass: probe_sglang now runs one synthetic image (grid 1x32x32 patches, HF + pixel_values layout) through model.visual under the device profiler and records it as the profile_run_kernels table + (same table vLLM''s own profile_run yields) - serving shows Triton `_fwd_kernel` x24 (= depth) for the vision attention, + the collector''s triton lane matches, gate ALIGNED. Every sglang/trtllm raw was re-probed afterwards (the probe file + is part of the execution fingerprint). sm90''s earlier aligned verdict for this gate rests on a record whose provenance + this checkout cannot establish; re-grade it with the new probe when the H20 box is next used. COLLECTOR defects found + by the smoke/gates that Hopper never reached (fa3/trtllm_mha paths don''t read them): MockModelRunner lacked prefill/decode_attention_backend_str + (flashinfer lane), the vision Triton kernel moved to sglang.kernels.ops.attention.prefill_attention, TritonAttnBackend + probes mambaish_config(model_config) so the MLA mock needs hf_config / linear_attn_registry_result, and the Gemma4 / + Llama4 case maps had no sm89 entry so the collector measured flashinfer while serving runs triton (arg_groups/model_overrides/gemma4.py, + llama4.py). All fixed with source citations and re-verified on hardware.' + applies_to: + - platform:sm89 + - framework:sglang-0.5.21 + - nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + - nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8 + - nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 + - nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + - Qwen/Qwen3.8-2.4T-A95B + - Qwen/Qwen3.8-2.4T-A95B-FP8 + - MiniMaxAI/MiniMax-M3 + - zai-org/GLM-5 + - instrument:collect_attn + - instrument:collect_mla + - instrument:collect_attn_encoder + sm89_trtllm_rc29_probe_2026_10_04: + statement: 'sm89 (L40) identity campaign, trtllm 1.3.0rc29 (digest 124da3f10b73). results/sm89/trtllm-1.3.0rc29.yaml: + 47 pass / 7 pass+custom / 70 fail (sm90: see results/sm90). HEADLINE: 26 cells (fp8-block 13, nvfp4 20 minus the 7 customized, + mxfp4 2, mxfp8 1) die at model init with ''no MoE implementation can serve this layer'': the rc29 template and typed + builder default moe_config.backend to CUTLASS and CutlassFusedMoE refuses FP8_BLOCK_SCALES below SM90/SM120 and NVFP4 + below SM100 - the same generator defect the B200 session found on sm100 (platform-blind CUTLASS default; hardware profile + `moe_backend` is keyed on hardware only). A/B on an empty L40 across the explicit backends (facts/ab_trt_moe_sm89, 12+ + probes): AUTO resolves to CUTLASS (same error, so a hardware-wide AUTO fact - tried and reverted - changes nothing); + TritonFusedMoE is SM90-only; Marlin implements nvfp4/w4a16_nvfp4 only (quant_unsupported for fp8_block, bf16, mxfp4); + VanillaMoE gets past selection and dies in CUDA-graph capture (cudaErrorStreamCaptureInvalidated). Re-verified in the + review pass (clean box, same evening): AUTO raws f883c2769dd6_AUTO (Qwen3-30B-A3B-FP8), 8610141a333f_AUTO (DeepSeek-R1), + 1ae319d4353f_AUTO (Nemotron-3-Nano-NVFP4), 56890dc997ec_AUTO (gpt-oss-20b MXFP4) all report ''requested CUTLASS'' and + the CUTLASS SM reason; DEEPGEMM, which serves fp8-block MoE on sm100, is turned down on Ada with ''DeepGemmFusedMoE + requires SM100, SM103, or SM107, got SM89'' (raws f883c2769dd6_DEEPGEMM, 8610141a333f_DEEPGEMM). So FP8-block / MXFP4 + MoE on rc29 is a PLATFORM FLOOR of the framework on Ada, while NVFP4 MoE is servable through MARLIN (weight-only). FIXED + IN THE GENERATOR (review pass): hardware.yaml l40s carries `moe_backend: {trtllm: {nvfp4: MARLIN}}` keyed by the artifact''s + quantization family (facts/apply.quant_family_of over utils._bundled_quantization), and the trtllm template now honours + the generic `moe_backend` param (it was silently dropped: the hardware fill ignored it and the template never read it; + a user value now beats the fill). The 10 roster NVFP4 MoE checkpoints re-probed under the plain render: 7 pass (MiniMax-M2.5/M2.7, + Nemotron-3 Nano/Super/Ultra, Qwen3.5-122B/397B), 3 reach later floors (Gemma-4: attentionOp ''Head size 512 is not supported + by MMHA'' native abort; Qwen3-235B: CutlassFp4GemmRunner needs Blackwell; Qwen3.8-2.4T capacity). MLA: the documented + ''no MLA FMHA below Hopper'' wall (collect_mla*.py FIXME(kernel-limit), ''Re-verify on the next framework version bump'') + was RE-VERIFIED on rc29 - with only the SM guard patched out a bf16 DeepSeek-V3 MLA module cell aborts in AttentionOp + ''Deepseek should be supported by fmha in generation part'' (attentionOp.cpp:3234, was :3091 on rc20; SIGSEGV in the + worker); guards kept, comments updated. Other fails: 8 DeepSeek-V4 ''requires Hopper or newer GPUs, got SM89'', 4 DeepGEMM + arch, 6 image gaps (transformers too old for DeepSeek-V4.1 / Muse-Glimmer; no `fla` for Kimi-K3 - same as sm100), 12 + generator rejects, 3 capacity, Gemma-4-31B-NVFP4 (flashinfer cute-DSL RMSNormFP4Quant ptxas needs .target sm_90), MiMo + arch not registered, Kimi-K2.5 quant parser (same as sm100). Gates: 7 of 9 real gates ALIGNED (attn ctx/gen + fp8-KV + gen, gdn ctx/gen, mamba2 ctx/gen), 13 declared sm89 floors (V3/R1 cannot boot: MoE + MLA; V3.2; V4). attn_ctx_fp8: serving''s + fp8-KV prefill runs fmha_v2 bf16 compute over the fp8 cache while the collector''s fp8-KV+fp8-context-FMHA case runs + the e4m3 kernel (fmha_v2_flash_attention_e4m3_..._sm89) - the gate recipe was moved to the bf16-compute combo that serving + selects; whether the SDK should bill the e4m3 FMHA on sm89 is an owner decision. encoder_attn was DIVERGED for lack + of evidence; CLOSED the same way as sglang (probe_trtllm runs the synthetic image through Qwen3VisionModel = mm_encoder.visual): + serving vision attention is `fmha_v2_flash_attention_bf16_64_32_S_qkv_64_sm89_kernel_nl` x24, the collector''s trtllm_mha + capture matches, ALIGNED. The two msa gate families the sm90 session left undeclared are declared with capture recipes + (trt_msa_ctx / trt_msa_gen) and GRADE FOR REAL on sm89 (ALIGNED) after the m3 dummy fix - see minimax_m3_dummy_dense_head_2026_10_05: + ''MiniMax-M3 does not boot on rc29'' was a dummy artifact, not a framework gap. On the other SMs the gate stays a FLOOR + whose note asks for a re-probe with the fixed dummy.' + applies_to: + - platform:sm89 + - framework:trtllm-1.3.0rc29 + - nvidia/DeepSeek-V3.1-NVFP4 + - deepseek-ai/DeepSeek-R1 + - deepseek-ai/DeepSeek-V3 + - Qwen/Qwen3-30B-A3B-FP8 + - openai/gpt-oss-120b + - openai/gpt-oss-20b + - nvidia/Gemma-4-26B-A4B-NVFP4 + - nvidia/Gemma-4-31B-IT-NVFP4 + - nvidia/Qwen3-235B-A22B-NVFP4 + - nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4 + - nvidia/MiniMax-M2.5-NVFP4 + - instrument:probe_driver + - instrument:collect_mla_module + - component:generator_facts + - MiniMaxAI/MiniMax-M2.7 + - Qwen/Qwen3-235B-A22B-FP8 + - XiaomiMiMo/MiMo-7B-Base + - XiaomiMiMo/MiMo-V2-Flash + sm89_collector_upgrade_vllm0300_sglang0521_trtllm_rc29_2026_10_04: + statement: 'COLLECTOR UPGRADE ON sm89 (L40; manifest pins vllm 0.30.0 / sglang 0.5.21 / trtllm 1.3.0rc29, images digest-pinned, + collect_campaign.build-image adds pyarrow/pandas): NOT a full data collection - plan audit, per-op --smoke (4 cases), + a 40-case --shuffle sample per op on all three frameworks (8 GPUs, one container per op, outputs under the probe workspace + collect/sample40), backend-truth gates (path_diff) and targeted reproductions. Plans (collect.py --model-cases-full + --sm 89 --plan-only): vllm 20 ops, sglang 27 ops, trtllm 23 ops; the DSA / DSV4-attention / GLM-5 sparse module ops + are floored at SM90 by cases/capabilities.yaml and plan to 0 cases on sm89. FIXES MADE (all source-cited, re-verified + on the L40): (1) capabilities.yaml op_min_sm 90 for dsv4_{csa,hca}_{context,generation}_module (was the only DSV4/GLM-5 + sparse family without a floor: every case queued and failed ''unexpected''); (2) sglang collect_attn MockModelRunner.prefill/decode_attention_backend_str + (flashinfer lane, sglang>=0.5.21); (3) sglang collect_attn_encoder: context_attention_fwd moved to sglang.kernels.ops.attention.prefill_attention; + (4) sglang collect_mla MockModelConfig.hf_config / linear_attn_registry_result / is_draft_model (TritonAttnBackend probes + mambaish_config); (5) Gemma4 and Llama4 case maps gain sm89: triton (arg_groups/model_overrides/gemma4.py, llama4.py) + - the collector measured flashinfer while serving runs triton; (6) trtllm MLA guards re-verified on rc29 (attentionOp.cpp:3234). + Items 2-4 are lanes Hopper (fa3) and Blackwell (trtllm_mha / trtllm_mla) never reach, so the sm90 0.5.21 port could + not have caught them; they will bite sm120 too (same flashinfer / triton lanes). SMOKE (4 cases/op): vllm 14 ops with + data, 6 floored; sglang 14 ops with data (mhc / dsv4 are deliberate UnverifiedCollector refusals); trtllm 13 ops with + data (MLA modules refuse by the verified guard). 40-CASE SAMPLE, failure classes (full table in HANDOFF_sm89.md): vllm + 0 failures in 11 of 14 ops, 1 capacity in attention_generation (96-head MHA, 4x4095 tokens = 24 GiB KV), moe 2 (1 capacity, + 1 block_n alignment), kda 5/12 (4 capacity, 1 partial grid); sglang: attention_context 13/40 = ONE deterministic class + - the plan schedules fp8_kv_cache_and_context_fmha for a FlashInfer SM (sm89/sm120 resolve flashinfer) and the collector + refuses it on purpose (''flashinfer has no FP8 prefill compute path'', a stale ''0.5.14'' in the message) - FIXED in + the review pass: get_context_attention_test_cases resolves the lane the same way run_attention_torch does (per-model + map, else the SM default table) and no longer plans fp8_context_fmha on flashinfer; sm89 context plan now 43,114 cases: + 18,752 bf16 + 18,752 fp8_kv_cache + 5,610 fp8_context_fmha on the triton/sink lanes (the 13,142 flashinfer-lane fp8_context_fmha + cases are gone), 4-case smoke clean; gdn 21/40 and kda 5/12 = the causal_conv1d Triton int32 token-offset overflow (total_tokens*conv_channels + >= 2**31, causal_conv1d_triton.py:373-379) at the largest grid points (reproduced: batch 32-64 x seq 16384-32768, 3 + of 112 points; the other 109 rows persist) - framework limit, not Ada; mla_generation 3 capacity + 1 IMA; trtllm: attention_context + 6 (4 CUDA faults + 2 Gemma4 trtllm-gen ''Blackwell-only'' guards), moe 8 (7 TP-shard alignment guards - the collector''s + own, known family - 1 capacity), gemm 2 stride-alignment RuntimeErrors, mhc 2 capacity. Capacity is the dominant Ada + difference: 44 GiB vs 80-141 GiB on Hopper moves the batch x kv corner of every attention / mla / msa / kda grid into + OOM, none of it a collector bug. NOT DONE: the full retained plan (hundreds of thousands of cases per framework - sm120''s + vllm run was ~12 h on 8 GPUs), perf-data finalize/publication, SDK consumer tests against new l40s data. DELIVERY BLOCKER + for any perf data from this box: the SDK system is l40s and its existing tables carry device ''NVIDIA L40S'' (checked + in the parquet device column), while this box is an 8x L40 (same AD102 die, lower TDP/clocks). Playbook section 8 forbids + relabelling a nearby product with the same SM, so L40 measurements must not be published under l40s; an `l40` system + definition would be needed first (the sm120 RTX PRO 5000 campaign hit the same wall).' + applies_to: + - platform:sm89 + - framework:vllm-0.30.0 + - framework:sglang-0.5.21 + - framework:trtllm-1.3.0rc29 + - instrument:collect.py + - instrument:collect_attn + - instrument:collect_mla + - instrument:collect_gdn + - instrument:collect_kda + - cases:capabilities.yaml + sm89_harness_lessons_2026_10_04: + statement: 'HARNESS LESSONS from the first Ada campaign (each fixed in code unless marked OPEN): (1) the SDK system has + no generator hardware profile -> resolve_facts raised and the pipeline swallowed it (l40s profile added, product-keyed + like h100/h200; same class as sm120); (2) golden-render cache keyed on git HEAD only: an uncommitted generator edit + re-rendered nothing (_generator_stamp = HEAD + hash of the dirty generator diff; B200 defect 3); (3) probes that die + by a native abort (trtllm attentionOp asserts, sglang argparse) wrote no raw and the matrix said only ''crashed before + dump'': the queue now keeps each probe''s console tail (archive/logs/.tail, 400 lines) and build_matrix classifies + from it (B200 defect 1); (4) tp>1 cells are emitted with the first N GPUs hard-coded (device=0,1,2,3) while other queues + run on those GPUs: sglang''s tp memory-balance check and trtllm''s executor then fail for reasons that are not the model''s. + Rerun alone on an idle group, the sglang Qwen3.8-2.4T cells show their real cause (CUDA OOM, capacity) and trtllm''s + ''Executor worker returned error'' still carries no cause in the raw (OPEN: the server probe does not surface the worker''s + message); a whole-group reservation in the queue plan is still needed (B200 defect 5, OPEN); (5) _fail_cause classified + on a 200-char slice of the last line: vllm''s selector error is cut before use_sparse=True and sglang/trtllm wrap the + real cause in epilogues, so the matrix filed DSA cells under a config cause and Nemotron under ''framework gap''. Now: + full last line first, then a deep pass over the whole traceback ONLY for unmatched causes; (6) FLOOR_SM took one SM, + and only the vllm verdict script had the floor mechanism (sglang / trtllm scripts now have it; comma lists everywhere, + workflow_check understands them); (7) 4 sglang and 4 trtllm gate recipes were never checked in (sgl_attn_ctx/gen/gemma4_hd512, + sgl_gdn_gen, trt_attn ctx/gen + fp8-KV) - written; (8) taxonomy: seeding sm89 from sm120 left 100+ kernel names unlabelled + on Ada (cuBLAS ampere_* GEMMs, flashinfer merge-states, FLA chunk kernels): collect/unlabeled.py style scan over raws + + captures is how the vocabulary was built and is worth promoting to a component option; (9) encoder_attn gates ask + phase profile_run, which only probe_vllm recorded: FIXED - probe_sglang / probe_trtllm run a synthetic image through + the vision module and record profile_run_kernels (sglang: inside the `trace` dict, where the record builder reads kernel + tables from; the first attempt wrote it top-level and was invisible); the probe file is part of the fingerprint, so + this costs a full sglang + trtllm re-probe (~2 h on 8 GPUs) and will on every SM that adopts it; (10) dummy staging + for gated repos needs preprocessor_config.json for Llama-4 (vllm builds the vision tower in its profile run) and the + tokenizer files: stage_tok.py in the workspace pulls them from unsloth / sibling mirrors - fetch_inputs.py''s bundled-config + fallback still stops short of that (OPEN); (11) workflow_check gates_declared is red for trtllm rc29 on EVERY SM (msa_context_module + / msa_generation_module have no trtllm gate; pre-existing on this branch, sm90 shows the same, OPEN); (12) tooling: + pkill -f / pgrep -f patterns match the invoking shell''s own command line - kill queue wrappers by pid. (13) `--emit-queues + --only X --plan-name P` REPLACES plan P with X''s runs only; build_records/build_matrix then resolve the dropped runs + from older plan files whose fingerprints no longer match the raws, and whole cells turn ''stale evidence'' (hit twice + on sm89). Either keep one plan file per --only set or re-emit the full set into it.' + applies_to: + - platform:sm89 + - instrument:probe_driver + - instrument:workflow_check + - instrument:fetch_inputs + - instrument:path_diff + - component:captures + minimax_m3_dummy_dense_head_2026_10_05: + statement: 'MINIMAX-M3 ON TRTLLM WAS NEVER A FRAMEWORK GAP - the harness dummy was. Diagnosis (L40, trtllm 1.3.0rc29, + 2026-10-05): a sitecustomize shim around MiniMaxM3SparseRuntimeBackend.forward showed the ONLY None argument is idx_k_cache; + MiniMaxM3KVCacheManagerV2.__init__ derives its sparse layer ids from explicit kwargs (no call site passes them), else + the executor''s sparse_attention_config (has no layer-id field), else the checkpoint convention `range(3, num_layers)` + (sparse/minimax_m3/cache_manager.py:275-303), while the model layer follows the config''s sparse_attention_freq. The + old m3 dummy cut kept two sparse/MoE layers renumbered 0..1 (freq [1, 1]) -> every layer sparse for the model, index-K + cache for no layer -> the ValueError the sm90 session (findings trtllm_sparse_config, hopper_trtllm_collector_pin_rc20_to_rc29_2026_10_02) + and the first sm89 pass both filed as a framework gap. Six config variants (no CUDA graph, no overlap scheduler, attn_backend + FLASHINFER, tokens_per_block 64/128) fail identically, because the cut, not the config, is wrong. FIX (dummies.variants_m3): + the sparse variant keeps the checkpoint''s 3-layer dense head in front of the first sparse layer (4 layers, freq [0, + 0, 0, 1]; same class as _ARCH_IMPLICIT_PERIODS - a framework constant the config cannot express). RESULT: MiniMaxAI/MiniMax-M3 + boots on trtllm rc29 on sm89 (kv auto and fp8), serving sparse attention = Triton `_block_max_score_kernel` + `_sparse_softmax_kernel` + (the MSA fmha_sm100 backend is SM100-only), collector trt_msa_ctx/gen captures match: msa_ctx/gen_MiniMax-M3 gates ALIGNED + on sm89; the msa FLOOR on the trtllm verdict script is now a re-probe request for sm90/sm100/sm103/sm120, not a framework + fact. SIDE EFFECT, vLLM: the all-sparse dummy had also been a FALSE PASS there - with a dense head present vllm 0.30.0 + dies at ''No common block size for 16'' (v1/worker/utils.py:386): the KV manager defaults to 16, FlashAttention accepts + MultipleOf(16), MiniMaxM3SparseBackend / MiniMaxM3IndexerBackend accept exactly [128] (models/minimax_m3/common/sparse_attention.py). + A/B on the L40: --block-size 128 boots (facts/ab_vllm_m3/bs128), 64 fails the same way. Generator fact added: models.yaml + profile `minimax-m3` with vllm block-size 128 (platform-neutral; the sm90 M3 vllm pass rests on the same all-sparse + dummy and the real 60-layer checkpoint has the dense head). sglang: with the 4-layer dummy MiniMax-M3 boots on sm89 + as well (pass+custom with the sm90 `cuda_graph_backend_prefill=breakable` customization; the bare default fails the + same mixed-chunk-prefill assertion as on sm90, facts/ab_sgl_m3/bare.json) - the earlier sm89 ''capacity'' verdict was + the 2-layer dummy''s. NVFP4 / MXFP8 M3 variants: platform floors on sm89 as before (trtllm Marlin: ''does not support + swiglu_gptoss_style''; sglang MXFP8 dense GEMM needs Blackwell).' + applies_to: + - platform:sm89 + - platform:sm90 + - framework:trtllm-1.3.0rc29 + - framework:vllm-0.30.0 + - MiniMaxAI/MiniMax-M3 + - MiniMaxAI/MiniMax-M3-MXFP8 + - nvidia/MiniMax-M3-NVFP4 + - instrument:dummies + - component:generator_facts + - family:m3 diff --git a/benchmarks/evidence/collector/op-probe-2026-10/results/retests/sm100/sglang-0.5.21.moe_auto.yaml b/benchmarks/evidence/collector/op-probe-2026-10/results/retests/sm100/sglang-0.5.21.moe_auto.yaml new file mode 100644 index 0000000000..0ecf6dd91b --- /dev/null +++ b/benchmarks/evidence/collector/op-probe-2026-10/results/retests/sm100/sglang-0.5.21.moe_auto.yaml @@ -0,0 +1,78 @@ +# Excerpt (2026-10-30): _meta plus the models collector comments cite; _meta.summary +# counts the full matrix, which lands with the op-probe harness. +_meta: + platform: b200_sm100 + framework: sglang + version: 0.5.21 + what: the MoE cells of results/sm100/-.yaml re-probed with the generator's hardware moe_backend fact removed from the rendered config; the as-rendered cells fail with cause 'generator fact rejected + by framework'. This file is the identity the framework picks on sm100 when it is allowed to choose. + override: --moe-runner-backend deepep_moe (hardware.yaml b200) dropped -> 0.5.21 AUTO + summary: + pass+override: 72 + fail: 3 + evidence: archive/retests/moe_auto/{raw,*.engine.yaml} in the B200 workspace (umbriel-b200-040, 2026-10-04) +results: + deepseek-ai/DeepSeek-V4-Flash: + verdict: pass+override + attention: DeepseekV4AttnBackend + moe_runner: FLASHINFER_MXFP4 + moe_kernel_families: + - trtllm_gen_moe + moe_kernels_seen: 4 + topology: tp1 + deepseek-ai/DeepSeek-V4-Flash-0731: + verdict: pass+override + attention: DeepseekV4AttnBackend + moe_runner: FLASHINFER_MXFP4 + moe_kernel_families: + - trtllm_gen_moe + moe_kernels_seen: 4 + topology: tp1 + deepseek-ai/DeepSeek-V4-Pro: + verdict: pass+override + attention: DeepseekV4AttnBackend + moe_runner: FLASHINFER_MXFP4 + moe_kernel_families: + - trtllm_gen_moe + moe_kernels_seen: 4 + topology: tp1 + deepseek-ai/DeepSeek-V4-Pro-0813: + verdict: pass+override + attention: DeepseekV4AttnBackend + moe_runner: FLASHINFER_MXFP4 + moe_kernel_families: + - trtllm_gen_moe + moe_kernels_seen: 4 + topology: tp1 + deepseek-ai/DeepSeek-V4.1-Flash: + verdict: pass+override + attention: DeepseekV4AttnBackend + moe_runner: FLASHINFER_MXFP4 + moe_kernel_families: + - trtllm_gen_moe + moe_kernels_seen: 4 + topology: tp1 + moonshotai/Kimi-K2.5: + verdict: pass+override + attention: TRTLLMMLABackend + moe_runner: FLASHINFER_TRTLLM + moe_kernel_families: + - trtllm_gen_moe + moe_kernels_seen: 5 + topology: tp1 + openai/gpt-oss-120b: + verdict: pass+override + attention: TRTLLMHAAttnBackend + moe_runner: FLASHINFER_MXFP4 + moe_kernel_families: + - trtllm_gen_moe + moe_kernels_seen: 5 + topology: tp1 + openai/gpt-oss-20b: + verdict: pass+override + attention: TRTLLMHAAttnBackend + moe_runner: FLASHINFER_MXFP4 + moe_kernel_families: + - trtllm_gen_moe + moe_kernels_seen: 5 + topology: tp1 diff --git a/benchmarks/evidence/collector/op-probe-2026-10/results/sm100/sglang-0.5.21.yaml b/benchmarks/evidence/collector/op-probe-2026-10/results/sm100/sglang-0.5.21.yaml new file mode 100644 index 0000000000..9f4188b3cb --- /dev/null +++ b/benchmarks/evidence/collector/op-probe-2026-10/results/sm100/sglang-0.5.21.yaml @@ -0,0 +1,34 @@ +# Excerpt (2026-10-30): _meta plus the models collector comments cite; _meta.summary +# counts the full matrix, which lands with the op-probe harness. +_meta: + platform: b200_sm100 + framework: sglang + version: 0.5.21 + verdicts: pass = plain `cli generate` output boots+runs; pass+custom = boots with facts-derived extra generate args; fail = root-caused (see results/findings.yaml) + summary: + pass: 102 + pass+custom: 6 + fail: 16 + golden_facts_not_applied: 0 +results: + moonshotai/Kimi-K2.5: + verdict: pass + evidence: current + attention: TRTLLMMLABackend + moe: CompressedTensorsFusedMoE->trtllm_gen_moe + kv_allocated: auto + topology: tp1 + openai/gpt-oss-120b: + verdict: pass + evidence: current + attention: TRTLLMHAAttnBackend + moe: Mxfp4MoE->trtllm_gen_moe + kv_allocated: auto + topology: tp1 + openai/gpt-oss-20b: + verdict: pass + evidence: current + attention: TRTLLMHAAttnBackend + moe: Mxfp4MoE->trtllm_gen_moe + kv_allocated: auto + topology: tp1 diff --git a/benchmarks/evidence/collector/op-probe-2026-10/results/sm100/trtllm-1.3.0rc29.yaml b/benchmarks/evidence/collector/op-probe-2026-10/results/sm100/trtllm-1.3.0rc29.yaml new file mode 100644 index 0000000000..fb3e2ed530 --- /dev/null +++ b/benchmarks/evidence/collector/op-probe-2026-10/results/sm100/trtllm-1.3.0rc29.yaml @@ -0,0 +1,17 @@ +# Excerpt (2026-10-30): _meta plus the models collector comments cite; _meta.summary +# counts the full matrix, which lands with the op-probe harness. +_meta: + platform: b200_sm100 + framework: trtllm + version: 1.3.0rc29 + verdicts: pass = plain `cli generate` output boots+runs; pass+custom = boots with facts-derived extra generate args; fail = root-caused (see results/findings.yaml) + summary: + pass: 78 + pass+custom: 0 + fail: 46 + golden_facts_not_applied: 0 +results: + moonshotai/Kimi-K2.5: + verdict: fail + cause: quant parser gap + error: 'TypeError: ''NoneType'' object is not subscriptable' diff --git a/benchmarks/evidence/collector/pr244/README.md b/benchmarks/evidence/collector/pr244/README.md index 76be4c224b..f0caba714e 100644 --- a/benchmarks/evidence/collector/pr244/README.md +++ b/benchmarks/evidence/collector/pr244/README.md @@ -6,8 +6,17 @@ SPDX-License-Identifier: Apache-2.0 # Collection audit evidence This non-packaged bundle verifies the retained vLLM 0.25.0 collection outcomes -and the 85 shipped Parquet tables. It contains case IDs, status vectors, -checksums, and table metadata, not diagnostic logs or machine paths. +and the 80 Parquet tables still shipped from it (85 originally). It contains case +IDs, status vectors, checksums, and table metadata, not diagnostic logs or +machine paths. + +The five `dsv4_hca_attn_module_perf` tables (one per system) were removed with +the hca_attn sidecar's SDK reader in #649. The case-outcome ledger keeps their +reported hashes, which record what was collected. The audit's `RETIRED_TABLES` +excludes them from the inventory check. The manifest was regenerated: it +drops their entries and carries the new hashes of the five 0.25.0 +`collection_meta.yaml` sidecars that listed them. Every other table entry is +unchanged. Run from the repository root in the development environment: diff --git a/benchmarks/evidence/collector/pr244/manifest.json b/benchmarks/evidence/collector/pr244/manifest.json index 960572fcfa..190122b1f9 100644 --- a/benchmarks/evidence/collector/pr244/manifest.json +++ b/benchmarks/evidence/collector/pr244/manifest.json @@ -253,8 +253,8 @@ "failed": 10739, "unattempted": 0 }, - "rows": 420111, - "table_count": 17, + "rows": 416193, + "table_count": 16, "tables": { "context_attention_perf.parquet": { "anomalies": { @@ -569,7 +569,7 @@ ], "rows": 5840, "sha256": "63e9b6b89418403e1998e691e2f93a5be67fd930a7997dac0d248ba0f577d7c3", - "sidecar_sha256": "bebbf880c0df4effe3b4d384e67984ddd36e254a83de2f0504e792646d6c8061" + "sidecar_sha256": "1ec714136ae98e456690f8773a44feba01be9e9a6f89f549b312a57596174af3" }, "dsv4_csa_generation_module_perf.parquet": { "anomalies": { @@ -731,177 +731,7 @@ ], "rows": 1544, "sha256": "632a9063ce5f857ab95731449176045fbde3551efcaefb35a93feea29549e5da", - "sidecar_sha256": "bebbf880c0df4effe3b4d384e67984ddd36e254a83de2f0504e792646d6c8061" - }, - "dsv4_hca_attn_module_perf.parquet": { - "anomalies": { - "duplicate_physical_keys": 0, - "nonfinite_latency": 0, - "nonpositive_latency": 0, - "null_cells": 0 - }, - "columns": { - "architecture": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "DeepseekV4ForCausalLM" - ] - }, - "batch_size": { - "distinct": 11, - "max": 1024, - "min": 1, - "nulls": 0, - "type": "int64" - }, - "compress_ratio": { - "distinct": 1, - "max": 128, - "min": 128, - "nulls": 0, - "type": "int64" - }, - "device": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "NVIDIA B300 SXM6 AC" - ] - }, - "framework": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "VLLM" - ] - }, - "gemm_type": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "fp8_block" - ] - }, - "isl": { - "distinct": 16, - "max": 8192, - "min": 1, - "nulls": 0, - "type": "int64" - }, - "kernel_source": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "FLASHMLA_SPARSE_DSV4" - ] - }, - "kv_cache_dtype": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "fp8" - ] - }, - "latency": { - "distinct": 1955, - "max": 5.2688, - "min": 0.0086, - "nulls": 0, - "type": "double" - }, - "local_num_heads": { - "distinct": 2, - "max": 128, - "min": 64, - "nulls": 0, - "type": "int64" - }, - "mla_dtype": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "bfloat16" - ] - }, - "model": { - "distinct": 2, - "nulls": 0, - "type": "string", - "values": [ - "sgl-project/DeepSeek-V4-Flash-FP8", - "sgl-project/DeepSeek-V4-Pro-FP8" - ] - }, - "num_heads": { - "distinct": 2, - "max": 128, - "min": 64, - "nulls": 0, - "type": "int64" - }, - "op_name": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "dsv4_hca_attn_module" - ] - }, - "step": { - "distinct": 26, - "max": 1048575, - "min": 0, - "nulls": 0, - "type": "int64" - }, - "tp_size": { - "distinct": 1, - "max": 1, - "min": 1, - "nulls": 0, - "type": "int64" - }, - "version": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "0.25.0" - ] - } - }, - "latency_unit": "ms", - "path": "python/aisimulate/src/aisimulate_core/systems/data/b300_sxm/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet", - "physical_key_columns": [ - "framework", - "version", - "op_name", - "kernel_source", - "model", - "architecture", - "mla_dtype", - "kv_cache_dtype", - "gemm_type", - "num_heads", - "local_num_heads", - "batch_size", - "isl", - "tp_size", - "step", - "compress_ratio" - ], - "rows": 3918, - "sha256": "cc616dbae378df19da3788d0d0b632b220144a60fb6067bbcb95cc94afe79513", - "sidecar_sha256": "bebbf880c0df4effe3b4d384e67984ddd36e254a83de2f0504e792646d6c8061" + "sidecar_sha256": "1ec714136ae98e456690f8773a44feba01be9e9a6f89f549b312a57596174af3" }, "dsv4_hca_context_module_perf.parquet": { "anomalies": { @@ -1063,7 +893,7 @@ ], "rows": 5840, "sha256": "fe349c63ed6445f442f2298e47456c37ffef5f1c43d458cf5d57538d282e581e", - "sidecar_sha256": "bebbf880c0df4effe3b4d384e67984ddd36e254a83de2f0504e792646d6c8061" + "sidecar_sha256": "1ec714136ae98e456690f8773a44feba01be9e9a6f89f549b312a57596174af3" }, "dsv4_hca_generation_module_perf.parquet": { "anomalies": { @@ -1225,7 +1055,7 @@ ], "rows": 1544, "sha256": "cdc48ea8a89e6b5543720df991f5f11d691b56349b1ab6a042bf6cbb550e6f62", - "sidecar_sha256": "bebbf880c0df4effe3b4d384e67984ddd36e254a83de2f0504e792646d6c8061" + "sidecar_sha256": "1ec714136ae98e456690f8773a44feba01be9e9a6f89f549b312a57596174af3" }, "dsv4_paged_mqa_logits_module_perf.parquet": { "anomalies": { @@ -1395,7 +1225,7 @@ ], "rows": 3918, "sha256": "64de2420ed80948ca068da0fc1e014734e53e6a77d71ab224e566b350be2280f", - "sidecar_sha256": "bebbf880c0df4effe3b4d384e67984ddd36e254a83de2f0504e792646d6c8061" + "sidecar_sha256": "1ec714136ae98e456690f8773a44feba01be9e9a6f89f549b312a57596174af3" }, "encoder_attention_perf.parquet": { "anomalies": { @@ -3034,8 +2864,8 @@ "failed": 10766, "unattempted": 0 }, - "rows": 420134, - "table_count": 17, + "rows": 416216, + "table_count": 16, "tables": { "context_attention_perf.parquet": { "anomalies": { @@ -3350,7 +3180,7 @@ ], "rows": 5840, "sha256": "76719d88e0481b5153c89ef0e1024435597f660de305bdb9f01750e9273189d3", - "sidecar_sha256": "ff4b25e84cb1c03a7c2b6bec9adc91e731a2dd57bb96d72089e13a4440d0d674" + "sidecar_sha256": "7fa57194fb7ad77faea184f07ef6c07ed293371aa4ac9429a9becf2fb2b4782d" }, "dsv4_csa_generation_module_perf.parquet": { "anomalies": { @@ -3512,9 +3342,9 @@ ], "rows": 1544, "sha256": "b1a09a8e004ebed907e1d6d85379c51f5ec4ca2c5812bb59a62327c99f2dc505", - "sidecar_sha256": "ff4b25e84cb1c03a7c2b6bec9adc91e731a2dd57bb96d72089e13a4440d0d674" + "sidecar_sha256": "7fa57194fb7ad77faea184f07ef6c07ed293371aa4ac9429a9becf2fb2b4782d" }, - "dsv4_hca_attn_module_perf.parquet": { + "dsv4_hca_context_module_perf.parquet": { "anomalies": { "duplicate_physical_keys": 0, "nonfinite_latency": 0, @@ -3569,8 +3399,8 @@ ] }, "isl": { - "distinct": 16, - "max": 8192, + "distinct": 21, + "max": 65536, "min": 1, "nulls": 0, "type": "int64" @@ -3592,19 +3422,12 @@ ] }, "latency": { - "distinct": 1985, - "max": 5.0301, - "min": 0.0106, + "distinct": 4931, + "max": 120.7941, + "min": 0.0436, "nulls": 0, "type": "double" }, - "local_num_heads": { - "distinct": 2, - "max": 128, - "min": 64, - "nulls": 0, - "type": "int64" - }, "mla_dtype": { "distinct": 1, "nulls": 0, @@ -3623,9 +3446,9 @@ ] }, "num_heads": { - "distinct": 2, + "distinct": 5, "max": 128, - "min": 64, + "min": 8, "nulls": 0, "type": "int64" }, @@ -3634,19 +3457,19 @@ "nulls": 0, "type": "string", "values": [ - "dsv4_hca_attn_module" + "dsv4_hca_context_module" ] }, "step": { - "distinct": 26, - "max": 1048575, + "distinct": 24, + "max": 65535, "min": 0, "nulls": 0, "type": "int64" }, "tp_size": { - "distinct": 1, - "max": 1, + "distinct": 4, + "max": 8, "min": 1, "nulls": 0, "type": "int64" @@ -3661,7 +3484,7 @@ } }, "latency_unit": "ms", - "path": "python/aisimulate/src/aisimulate_core/systems/data/gb200/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet", + "path": "python/aisimulate/src/aisimulate_core/systems/data/gb200/sparse_attention/vllm/0.25.0/dsv4_hca_context_module_perf.parquet", "physical_key_columns": [ "framework", "version", @@ -3673,18 +3496,17 @@ "kv_cache_dtype", "gemm_type", "num_heads", - "local_num_heads", "batch_size", "isl", "tp_size", "step", "compress_ratio" ], - "rows": 3918, - "sha256": "0abf339525484c011c1c57508bbf01d1338f98f93cbc3d0debc6b2b383b2d0dc", - "sidecar_sha256": "ff4b25e84cb1c03a7c2b6bec9adc91e731a2dd57bb96d72089e13a4440d0d674" + "rows": 5840, + "sha256": "633e6a30fc69935c0edd9fb672d26a02ba3d13a36e5412f5dab57571b37c7ec1", + "sidecar_sha256": "7fa57194fb7ad77faea184f07ef6c07ed293371aa4ac9429a9becf2fb2b4782d" }, - "dsv4_hca_context_module_perf.parquet": { + "dsv4_hca_generation_module_perf.parquet": { "anomalies": { "duplicate_physical_keys": 0, "nonfinite_latency": 0, @@ -3739,8 +3561,8 @@ ] }, "isl": { - "distinct": 21, - "max": 65536, + "distinct": 1, + "max": 1, "min": 1, "nulls": 0, "type": "int64" @@ -3762,9 +3584,9 @@ ] }, "latency": { - "distinct": 4931, - "max": 120.7941, - "min": 0.0436, + "distinct": 740, + "max": 0.5519, + "min": 0.0421, "nulls": 0, "type": "double" }, @@ -3797,13 +3619,13 @@ "nulls": 0, "type": "string", "values": [ - "dsv4_hca_context_module" + "dsv4_hca_generation_module" ] }, "step": { - "distinct": 24, - "max": 65535, - "min": 0, + "distinct": 21, + "max": 65536, + "min": 1, "nulls": 0, "type": "int64" }, @@ -3824,7 +3646,7 @@ } }, "latency_unit": "ms", - "path": "python/aisimulate/src/aisimulate_core/systems/data/gb200/sparse_attention/vllm/0.25.0/dsv4_hca_context_module_perf.parquet", + "path": "python/aisimulate/src/aisimulate_core/systems/data/gb200/sparse_attention/vllm/0.25.0/dsv4_hca_generation_module_perf.parquet", "physical_key_columns": [ "framework", "version", @@ -3842,11 +3664,11 @@ "step", "compress_ratio" ], - "rows": 5840, - "sha256": "633e6a30fc69935c0edd9fb672d26a02ba3d13a36e5412f5dab57571b37c7ec1", - "sidecar_sha256": "ff4b25e84cb1c03a7c2b6bec9adc91e731a2dd57bb96d72089e13a4440d0d674" + "rows": 1544, + "sha256": "9807d59d5b0e6a14a5d82feb7b81c8529fbe2d0d65120fd32706a7b68fe4cf2a", + "sidecar_sha256": "7fa57194fb7ad77faea184f07ef6c07ed293371aa4ac9429a9becf2fb2b4782d" }, - "dsv4_hca_generation_module_perf.parquet": { + "dsv4_paged_mqa_logits_module_perf.parquet": { "anomalies": { "duplicate_physical_keys": 0, "nonfinite_latency": 0, @@ -3871,8 +3693,8 @@ }, "compress_ratio": { "distinct": 1, - "max": 128, - "min": 128, + "max": 4, + "min": 4, "nulls": 0, "type": "int64" }, @@ -3901,8 +3723,8 @@ ] }, "isl": { - "distinct": 1, - "max": 1, + "distinct": 16, + "max": 8192, "min": 1, "nulls": 0, "type": "int64" @@ -3912,7 +3734,7 @@ "nulls": 0, "type": "string", "values": [ - "FLASHMLA_SPARSE_DSV4" + "vllm.utils.deep_gemm.fp8_fp4_paged_mqa_logits" ] }, "kv_cache_dtype": { @@ -3924,18 +3746,25 @@ ] }, "latency": { - "distinct": 740, - "max": 0.5519, - "min": 0.0421, + "distinct": 992, + "max": 11.271, + "min": 0.0062, "nulls": 0, "type": "double" }, + "local_num_heads": { + "distinct": 2, + "max": 128, + "min": 64, + "nulls": 0, + "type": "int64" + }, "mla_dtype": { "distinct": 1, "nulls": 0, "type": "string", "values": [ - "bfloat16" + "fp8_e4m3" ] }, "model": { @@ -3948,9 +3777,9 @@ ] }, "num_heads": { - "distinct": 5, + "distinct": 2, "max": 128, - "min": 8, + "min": 64, "nulls": 0, "type": "int64" }, @@ -3959,176 +3788,7 @@ "nulls": 0, "type": "string", "values": [ - "dsv4_hca_generation_module" - ] - }, - "step": { - "distinct": 21, - "max": 65536, - "min": 1, - "nulls": 0, - "type": "int64" - }, - "tp_size": { - "distinct": 4, - "max": 8, - "min": 1, - "nulls": 0, - "type": "int64" - }, - "version": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "0.25.0" - ] - } - }, - "latency_unit": "ms", - "path": "python/aisimulate/src/aisimulate_core/systems/data/gb200/sparse_attention/vllm/0.25.0/dsv4_hca_generation_module_perf.parquet", - "physical_key_columns": [ - "framework", - "version", - "op_name", - "kernel_source", - "model", - "architecture", - "mla_dtype", - "kv_cache_dtype", - "gemm_type", - "num_heads", - "batch_size", - "isl", - "tp_size", - "step", - "compress_ratio" - ], - "rows": 1544, - "sha256": "9807d59d5b0e6a14a5d82feb7b81c8529fbe2d0d65120fd32706a7b68fe4cf2a", - "sidecar_sha256": "ff4b25e84cb1c03a7c2b6bec9adc91e731a2dd57bb96d72089e13a4440d0d674" - }, - "dsv4_paged_mqa_logits_module_perf.parquet": { - "anomalies": { - "duplicate_physical_keys": 0, - "nonfinite_latency": 0, - "nonpositive_latency": 0, - "null_cells": 0 - }, - "columns": { - "architecture": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "DeepseekV4ForCausalLM" - ] - }, - "batch_size": { - "distinct": 11, - "max": 1024, - "min": 1, - "nulls": 0, - "type": "int64" - }, - "compress_ratio": { - "distinct": 1, - "max": 4, - "min": 4, - "nulls": 0, - "type": "int64" - }, - "device": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "NVIDIA GB200" - ] - }, - "framework": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "VLLM" - ] - }, - "gemm_type": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "fp8_block" - ] - }, - "isl": { - "distinct": 16, - "max": 8192, - "min": 1, - "nulls": 0, - "type": "int64" - }, - "kernel_source": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "vllm.utils.deep_gemm.fp8_fp4_paged_mqa_logits" - ] - }, - "kv_cache_dtype": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "fp8" - ] - }, - "latency": { - "distinct": 992, - "max": 11.271, - "min": 0.0062, - "nulls": 0, - "type": "double" - }, - "local_num_heads": { - "distinct": 2, - "max": 128, - "min": 64, - "nulls": 0, - "type": "int64" - }, - "mla_dtype": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "fp8_e4m3" - ] - }, - "model": { - "distinct": 2, - "nulls": 0, - "type": "string", - "values": [ - "sgl-project/DeepSeek-V4-Flash-FP8", - "sgl-project/DeepSeek-V4-Pro-FP8" - ] - }, - "num_heads": { - "distinct": 2, - "max": 128, - "min": 64, - "nulls": 0, - "type": "int64" - }, - "op_name": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "dsv4_paged_mqa_logits_module" + "dsv4_paged_mqa_logits_module" ] }, "step": { @@ -4176,7 +3836,7 @@ ], "rows": 3918, "sha256": "e62d12cc3c8a125b66ad066de2abfda101c06a6acdbf9f7d0e9bebd4f442f4c3", - "sidecar_sha256": "ff4b25e84cb1c03a7c2b6bec9adc91e731a2dd57bb96d72089e13a4440d0d674" + "sidecar_sha256": "7fa57194fb7ad77faea184f07ef6c07ed293371aa4ac9429a9becf2fb2b4782d" }, "encoder_attention_perf.parquet": { "anomalies": { @@ -5815,8 +5475,8 @@ "failed": 10699, "unattempted": 0 }, - "rows": 420226, - "table_count": 17, + "rows": 416308, + "table_count": 16, "tables": { "context_attention_perf.parquet": { "anomalies": { @@ -6131,7 +5791,7 @@ ], "rows": 5840, "sha256": "81276eea4e424b2146b58256cd1697529a3a00c8bf3ee7c6613d54d6094972fc", - "sidecar_sha256": "ff4b25e84cb1c03a7c2b6bec9adc91e731a2dd57bb96d72089e13a4440d0d674" + "sidecar_sha256": "7fa57194fb7ad77faea184f07ef6c07ed293371aa4ac9429a9becf2fb2b4782d" }, "dsv4_csa_generation_module_perf.parquet": { "anomalies": { @@ -6293,177 +5953,7 @@ ], "rows": 1544, "sha256": "eb54cfe271dc3d15d0849745c17cb1979da4ac4afc63504acece52bcf4482d08", - "sidecar_sha256": "ff4b25e84cb1c03a7c2b6bec9adc91e731a2dd57bb96d72089e13a4440d0d674" - }, - "dsv4_hca_attn_module_perf.parquet": { - "anomalies": { - "duplicate_physical_keys": 0, - "nonfinite_latency": 0, - "nonpositive_latency": 0, - "null_cells": 0 - }, - "columns": { - "architecture": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "DeepseekV4ForCausalLM" - ] - }, - "batch_size": { - "distinct": 11, - "max": 1024, - "min": 1, - "nulls": 0, - "type": "int64" - }, - "compress_ratio": { - "distinct": 1, - "max": 128, - "min": 128, - "nulls": 0, - "type": "int64" - }, - "device": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "NVIDIA GB300" - ] - }, - "framework": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "VLLM" - ] - }, - "gemm_type": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "fp8_block" - ] - }, - "isl": { - "distinct": 16, - "max": 8192, - "min": 1, - "nulls": 0, - "type": "int64" - }, - "kernel_source": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "FLASHMLA_SPARSE_DSV4" - ] - }, - "kv_cache_dtype": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "fp8" - ] - }, - "latency": { - "distinct": 1920, - "max": 4.9961, - "min": 0.0106, - "nulls": 0, - "type": "double" - }, - "local_num_heads": { - "distinct": 2, - "max": 128, - "min": 64, - "nulls": 0, - "type": "int64" - }, - "mla_dtype": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "bfloat16" - ] - }, - "model": { - "distinct": 2, - "nulls": 0, - "type": "string", - "values": [ - "sgl-project/DeepSeek-V4-Flash-FP8", - "sgl-project/DeepSeek-V4-Pro-FP8" - ] - }, - "num_heads": { - "distinct": 2, - "max": 128, - "min": 64, - "nulls": 0, - "type": "int64" - }, - "op_name": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "dsv4_hca_attn_module" - ] - }, - "step": { - "distinct": 26, - "max": 1048575, - "min": 0, - "nulls": 0, - "type": "int64" - }, - "tp_size": { - "distinct": 1, - "max": 1, - "min": 1, - "nulls": 0, - "type": "int64" - }, - "version": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "0.25.0" - ] - } - }, - "latency_unit": "ms", - "path": "python/aisimulate/src/aisimulate_core/systems/data/gb300/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet", - "physical_key_columns": [ - "framework", - "version", - "op_name", - "kernel_source", - "model", - "architecture", - "mla_dtype", - "kv_cache_dtype", - "gemm_type", - "num_heads", - "local_num_heads", - "batch_size", - "isl", - "tp_size", - "step", - "compress_ratio" - ], - "rows": 3918, - "sha256": "49ecc0d7e34e1503f7bdbea1c3599bb20a219488786a22fef43cf216de641383", - "sidecar_sha256": "ff4b25e84cb1c03a7c2b6bec9adc91e731a2dd57bb96d72089e13a4440d0d674" + "sidecar_sha256": "7fa57194fb7ad77faea184f07ef6c07ed293371aa4ac9429a9becf2fb2b4782d" }, "dsv4_hca_context_module_perf.parquet": { "anomalies": { @@ -6625,7 +6115,7 @@ ], "rows": 5840, "sha256": "56d026f2e2b5b843c0afbe6b6c0d9b3489abf5db256d5d61ee065f7bcd72d15f", - "sidecar_sha256": "ff4b25e84cb1c03a7c2b6bec9adc91e731a2dd57bb96d72089e13a4440d0d674" + "sidecar_sha256": "7fa57194fb7ad77faea184f07ef6c07ed293371aa4ac9429a9becf2fb2b4782d" }, "dsv4_hca_generation_module_perf.parquet": { "anomalies": { @@ -6787,7 +6277,7 @@ ], "rows": 1544, "sha256": "3dee3c879fe7fdd9fd1e111d2f8df79cf15d8e635ebe32f123b94adc441fc31e", - "sidecar_sha256": "ff4b25e84cb1c03a7c2b6bec9adc91e731a2dd57bb96d72089e13a4440d0d674" + "sidecar_sha256": "7fa57194fb7ad77faea184f07ef6c07ed293371aa4ac9429a9becf2fb2b4782d" }, "dsv4_paged_mqa_logits_module_perf.parquet": { "anomalies": { @@ -6957,7 +6447,7 @@ ], "rows": 3918, "sha256": "12f3aa8b290193bf97e2f144c9bdda347eeba8ab10698669bc41d1fa617280fd", - "sidecar_sha256": "ff4b25e84cb1c03a7c2b6bec9adc91e731a2dd57bb96d72089e13a4440d0d674" + "sidecar_sha256": "7fa57194fb7ad77faea184f07ef6c07ed293371aa4ac9429a9becf2fb2b4782d" }, "encoder_attention_perf.parquet": { "anomalies": { @@ -8586,8 +8076,8 @@ "failed": 3495, "unattempted": 0 }, - "rows": 353449, - "table_count": 17, + "rows": 349531, + "table_count": 16, "tables": { "context_attention_perf.parquet": { "anomalies": { @@ -8710,200 +8200,38 @@ "values": [ "0.25.0" ] - }, - "window_size": { - "distinct": 6, - "max": 8192, - "min": 0, - "nulls": 0, - "type": "int64" - } - }, - "latency_unit": "ms", - "path": "python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/attention/vllm/0.25.0/context_attention_perf.parquet", - "physical_key_columns": [ - "framework", - "version", - "op_name", - "kernel_source", - "batch_size", - "isl", - "num_heads", - "num_key_value_heads", - "head_dim", - "beam_width", - "attn_dtype", - "kv_cache_dtype", - "step", - "window_size" - ], - "rows": 53771, - "sha256": "7bca62aa7d140efcd54dadb60ff7582de130f7693265c26a53a817e92a601f9a", - "sidecar_sha256": "9c22cf1a4cb9d50fad792d7e319e933b10cf57feae22c0fb90a46127ce12fc46" - }, - "dsv4_csa_context_module_perf.parquet": { - "anomalies": { - "duplicate_physical_keys": 0, - "nonfinite_latency": 0, - "nonpositive_latency": 0, - "null_cells": 0 - }, - "columns": { - "architecture": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "DeepseekV4ForCausalLM" - ] - }, - "batch_size": { - "distinct": 11, - "max": 1024, - "min": 1, - "nulls": 0, - "type": "int64" - }, - "compress_ratio": { - "distinct": 1, - "max": 4, - "min": 4, - "nulls": 0, - "type": "int64" - }, - "device": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "NVIDIA H100 80GB HBM3" - ] - }, - "framework": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "VLLM" - ] - }, - "gemm_type": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "fp8_block" - ] - }, - "isl": { - "distinct": 21, - "max": 65536, - "min": 1, - "nulls": 0, - "type": "int64" - }, - "kernel_source": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "FLASHMLA_SPARSE_DSV4" - ] - }, - "kv_cache_dtype": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "fp8" - ] - }, - "latency": { - "distinct": 5417, - "max": 357.5721, - "min": 0.0702, - "nulls": 0, - "type": "double" - }, - "mla_dtype": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "bfloat16" - ] - }, - "model": { - "distinct": 2, - "nulls": 0, - "type": "string", - "values": [ - "sgl-project/DeepSeek-V4-Flash-FP8", - "sgl-project/DeepSeek-V4-Pro-FP8" - ] - }, - "num_heads": { - "distinct": 5, - "max": 128, - "min": 8, - "nulls": 0, - "type": "int64" - }, - "op_name": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "dsv4_csa_context_module" - ] - }, - "step": { - "distinct": 24, - "max": 65535, - "min": 0, - "nulls": 0, - "type": "int64" - }, - "tp_size": { - "distinct": 4, - "max": 8, - "min": 1, - "nulls": 0, - "type": "int64" - }, - "version": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "0.25.0" - ] + }, + "window_size": { + "distinct": 6, + "max": 8192, + "min": 0, + "nulls": 0, + "type": "int64" } }, "latency_unit": "ms", - "path": "python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.25.0/dsv4_csa_context_module_perf.parquet", + "path": "python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/attention/vllm/0.25.0/context_attention_perf.parquet", "physical_key_columns": [ "framework", "version", "op_name", "kernel_source", - "model", - "architecture", - "mla_dtype", - "kv_cache_dtype", - "gemm_type", - "num_heads", "batch_size", "isl", - "tp_size", + "num_heads", + "num_key_value_heads", + "head_dim", + "beam_width", + "attn_dtype", + "kv_cache_dtype", "step", - "compress_ratio" + "window_size" ], - "rows": 5789, - "sha256": "fd4664e02b104be261b238db5dd571544dc35018ffa6e5fdda700f68e0503a79", - "sidecar_sha256": "19017dfd3c09a7df0dc5dd4d47e01ff0d9324998f3fe4f93b999effe36bfe8aa" + "rows": 53771, + "sha256": "7bca62aa7d140efcd54dadb60ff7582de130f7693265c26a53a817e92a601f9a", + "sidecar_sha256": "9c22cf1a4cb9d50fad792d7e319e933b10cf57feae22c0fb90a46127ce12fc46" }, - "dsv4_csa_generation_module_perf.parquet": { + "dsv4_csa_context_module_perf.parquet": { "anomalies": { "duplicate_physical_keys": 0, "nonfinite_latency": 0, @@ -8958,8 +8286,8 @@ ] }, "isl": { - "distinct": 1, - "max": 1, + "distinct": 21, + "max": 65536, "min": 1, "nulls": 0, "type": "int64" @@ -8981,9 +8309,9 @@ ] }, "latency": { - "distinct": 1031, - "max": 1.3585, - "min": 0.0695, + "distinct": 5417, + "max": 357.5721, + "min": 0.0702, "nulls": 0, "type": "double" }, @@ -9016,13 +8344,13 @@ "nulls": 0, "type": "string", "values": [ - "dsv4_csa_generation_module" + "dsv4_csa_context_module" ] }, "step": { - "distinct": 21, - "max": 65536, - "min": 1, + "distinct": 24, + "max": 65535, + "min": 0, "nulls": 0, "type": "int64" }, @@ -9043,7 +8371,7 @@ } }, "latency_unit": "ms", - "path": "python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.25.0/dsv4_csa_generation_module_perf.parquet", + "path": "python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.25.0/dsv4_csa_context_module_perf.parquet", "physical_key_columns": [ "framework", "version", @@ -9061,11 +8389,11 @@ "step", "compress_ratio" ], - "rows": 1544, - "sha256": "c9a6ff6f081dee79ff3053d32a8e85d550ee7aaab0c3d080bb2d33d7ca4ab839", - "sidecar_sha256": "19017dfd3c09a7df0dc5dd4d47e01ff0d9324998f3fe4f93b999effe36bfe8aa" + "rows": 5789, + "sha256": "fd4664e02b104be261b238db5dd571544dc35018ffa6e5fdda700f68e0503a79", + "sidecar_sha256": "c17befda9060286fc7bcbe28d086e8964c47644c700e1f60d604ccbed20d6486" }, - "dsv4_hca_attn_module_perf.parquet": { + "dsv4_csa_generation_module_perf.parquet": { "anomalies": { "duplicate_physical_keys": 0, "nonfinite_latency": 0, @@ -9090,8 +8418,8 @@ }, "compress_ratio": { "distinct": 1, - "max": 128, - "min": 128, + "max": 4, + "min": 4, "nulls": 0, "type": "int64" }, @@ -9120,8 +8448,8 @@ ] }, "isl": { - "distinct": 16, - "max": 8192, + "distinct": 1, + "max": 1, "min": 1, "nulls": 0, "type": "int64" @@ -9143,19 +8471,12 @@ ] }, "latency": { - "distinct": 2358, - "max": 14.2497, - "min": 0.01, + "distinct": 1031, + "max": 1.3585, + "min": 0.0695, "nulls": 0, "type": "double" }, - "local_num_heads": { - "distinct": 2, - "max": 128, - "min": 64, - "nulls": 0, - "type": "int64" - }, "mla_dtype": { "distinct": 1, "nulls": 0, @@ -9174,9 +8495,9 @@ ] }, "num_heads": { - "distinct": 2, + "distinct": 5, "max": 128, - "min": 64, + "min": 8, "nulls": 0, "type": "int64" }, @@ -9185,19 +8506,19 @@ "nulls": 0, "type": "string", "values": [ - "dsv4_hca_attn_module" + "dsv4_csa_generation_module" ] }, "step": { - "distinct": 26, - "max": 1048575, - "min": 0, + "distinct": 21, + "max": 65536, + "min": 1, "nulls": 0, "type": "int64" }, "tp_size": { - "distinct": 1, - "max": 1, + "distinct": 4, + "max": 8, "min": 1, "nulls": 0, "type": "int64" @@ -9212,7 +8533,7 @@ } }, "latency_unit": "ms", - "path": "python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet", + "path": "python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.25.0/dsv4_csa_generation_module_perf.parquet", "physical_key_columns": [ "framework", "version", @@ -9224,16 +8545,15 @@ "kv_cache_dtype", "gemm_type", "num_heads", - "local_num_heads", "batch_size", "isl", "tp_size", "step", "compress_ratio" ], - "rows": 3918, - "sha256": "86ca415600d3f4b468d683a689403ca298739cc65967c22c66fc2108ae6aa2fe", - "sidecar_sha256": "19017dfd3c09a7df0dc5dd4d47e01ff0d9324998f3fe4f93b999effe36bfe8aa" + "rows": 1544, + "sha256": "c9a6ff6f081dee79ff3053d32a8e85d550ee7aaab0c3d080bb2d33d7ca4ab839", + "sidecar_sha256": "c17befda9060286fc7bcbe28d086e8964c47644c700e1f60d604ccbed20d6486" }, "dsv4_hca_context_module_perf.parquet": { "anomalies": { @@ -9395,7 +8715,7 @@ ], "rows": 5793, "sha256": "9b0d69b9a9634ffd2d58b8e38d402db292ab003f5d4a59cab77a15c983f13d0d", - "sidecar_sha256": "19017dfd3c09a7df0dc5dd4d47e01ff0d9324998f3fe4f93b999effe36bfe8aa" + "sidecar_sha256": "c17befda9060286fc7bcbe28d086e8964c47644c700e1f60d604ccbed20d6486" }, "dsv4_hca_generation_module_perf.parquet": { "anomalies": { @@ -9557,7 +8877,7 @@ ], "rows": 1544, "sha256": "ddeed06756d6e93d3afa9d372587ed19edae41f6051435816b27fc265e7689a8", - "sidecar_sha256": "19017dfd3c09a7df0dc5dd4d47e01ff0d9324998f3fe4f93b999effe36bfe8aa" + "sidecar_sha256": "c17befda9060286fc7bcbe28d086e8964c47644c700e1f60d604ccbed20d6486" }, "dsv4_paged_mqa_logits_module_perf.parquet": { "anomalies": { @@ -9727,7 +9047,7 @@ ], "rows": 3918, "sha256": "8eaca4d81931d542e879f8e1cd9728dda466ecd5d9c19449eef7be5876e90270", - "sidecar_sha256": "19017dfd3c09a7df0dc5dd4d47e01ff0d9324998f3fe4f93b999effe36bfe8aa" + "sidecar_sha256": "c17befda9060286fc7bcbe28d086e8964c47644c700e1f60d604ccbed20d6486" }, "encoder_attention_perf.parquet": { "anomalies": { @@ -11343,8 +10663,8 @@ "failed": 3394, "unattempted": 0 }, - "rows": 353658, - "table_count": 17, + "rows": 349740, + "table_count": 16, "tables": { "context_attention_perf.parquet": { "anomalies": { @@ -11658,7 +10978,7 @@ ], "rows": 5840, "sha256": "8fec19f573b830fd854b0feb199cc6841f8070a86308dc892eec80afdaf93c29", - "sidecar_sha256": "23f9cb39768514c0ef6f079e9bc43abc2719cd464638e371fda85f2922ce7bc1" + "sidecar_sha256": "2c8891395cac294aadefdfe36aa81da7915eb5fb14a581f6a1cc5722957c6d43" }, "dsv4_csa_generation_module_perf.parquet": { "anomalies": { @@ -11820,177 +11140,7 @@ ], "rows": 1544, "sha256": "a4c1dc94f505ba0360125de51735948185b0e9bd44128a42668613875ae7741a", - "sidecar_sha256": "23f9cb39768514c0ef6f079e9bc43abc2719cd464638e371fda85f2922ce7bc1" - }, - "dsv4_hca_attn_module_perf.parquet": { - "anomalies": { - "duplicate_physical_keys": 0, - "nonfinite_latency": 0, - "nonpositive_latency": 0, - "null_cells": 0 - }, - "columns": { - "architecture": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "DeepseekV4ForCausalLM" - ] - }, - "batch_size": { - "distinct": 11, - "max": 1024, - "min": 1, - "nulls": 0, - "type": "int64" - }, - "compress_ratio": { - "distinct": 1, - "max": 128, - "min": 128, - "nulls": 0, - "type": "int64" - }, - "device": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "NVIDIA H200" - ] - }, - "framework": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "VLLM" - ] - }, - "gemm_type": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "fp8_block" - ] - }, - "isl": { - "distinct": 16, - "max": 8192, - "min": 1, - "nulls": 0, - "type": "int64" - }, - "kernel_source": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "FLASHMLA_SPARSE_DSV4" - ] - }, - "kv_cache_dtype": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "fp8" - ] - }, - "latency": { - "distinct": 2307, - "max": 14.4426, - "min": 0.0096, - "nulls": 0, - "type": "double" - }, - "local_num_heads": { - "distinct": 2, - "max": 128, - "min": 64, - "nulls": 0, - "type": "int64" - }, - "mla_dtype": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "bfloat16" - ] - }, - "model": { - "distinct": 2, - "nulls": 0, - "type": "string", - "values": [ - "sgl-project/DeepSeek-V4-Flash-FP8", - "sgl-project/DeepSeek-V4-Pro-FP8" - ] - }, - "num_heads": { - "distinct": 2, - "max": 128, - "min": 64, - "nulls": 0, - "type": "int64" - }, - "op_name": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "dsv4_hca_attn_module" - ] - }, - "step": { - "distinct": 26, - "max": 1048575, - "min": 0, - "nulls": 0, - "type": "int64" - }, - "tp_size": { - "distinct": 1, - "max": 1, - "min": 1, - "nulls": 0, - "type": "int64" - }, - "version": { - "distinct": 1, - "nulls": 0, - "type": "string", - "values": [ - "0.25.0" - ] - } - }, - "latency_unit": "ms", - "path": "python/aisimulate/src/aisimulate_core/systems/data/h200_sxm/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet", - "physical_key_columns": [ - "framework", - "version", - "op_name", - "kernel_source", - "model", - "architecture", - "mla_dtype", - "kv_cache_dtype", - "gemm_type", - "num_heads", - "local_num_heads", - "batch_size", - "isl", - "tp_size", - "step", - "compress_ratio" - ], - "rows": 3918, - "sha256": "3b0d87916ae3f6414cbe92966632b9133fb162a596ffce6b44cfeb916da3a4af", - "sidecar_sha256": "23f9cb39768514c0ef6f079e9bc43abc2719cd464638e371fda85f2922ce7bc1" + "sidecar_sha256": "2c8891395cac294aadefdfe36aa81da7915eb5fb14a581f6a1cc5722957c6d43" }, "dsv4_hca_context_module_perf.parquet": { "anomalies": { @@ -12152,7 +11302,7 @@ ], "rows": 5840, "sha256": "39fa711b357b5c72251862cc741ae81e6a486f6ad0adb2ae29e25565d99642b1", - "sidecar_sha256": "23f9cb39768514c0ef6f079e9bc43abc2719cd464638e371fda85f2922ce7bc1" + "sidecar_sha256": "2c8891395cac294aadefdfe36aa81da7915eb5fb14a581f6a1cc5722957c6d43" }, "dsv4_hca_generation_module_perf.parquet": { "anomalies": { @@ -12314,7 +11464,7 @@ ], "rows": 1544, "sha256": "d9683b5d8a43df2f023bfb4ac50d38f8717b0c447d2ed3962b3ae5eaa64fbf00", - "sidecar_sha256": "23f9cb39768514c0ef6f079e9bc43abc2719cd464638e371fda85f2922ce7bc1" + "sidecar_sha256": "2c8891395cac294aadefdfe36aa81da7915eb5fb14a581f6a1cc5722957c6d43" }, "dsv4_paged_mqa_logits_module_perf.parquet": { "anomalies": { @@ -12484,7 +11634,7 @@ ], "rows": 3918, "sha256": "ad63707a693a23b190807366e1c6eb65cacc33ef42a92b070c23f6985ed0514b", - "sidecar_sha256": "23f9cb39768514c0ef6f079e9bc43abc2719cd464638e371fda85f2922ce7bc1" + "sidecar_sha256": "2c8891395cac294aadefdfe36aa81da7915eb5fb14a581f6a1cc5722957c6d43" }, "encoder_attention_perf.parquet": { "anomalies": { diff --git a/crates/core/src/perfmodel/common/enums.rs b/crates/core/src/perfmodel/common/enums.rs index 9e024319fa..78a0f5c952 100644 --- a/crates/core/src/perfmodel/common/enums.rs +++ b/crates/core/src/perfmodel/common/enums.rs @@ -641,8 +641,6 @@ pub enum PerfDataFilename { Dsv4CsaGenerationModule, Dsv4HcaGenerationModule, Dsv4PagedMqaLogitsModule, - Dsv4HcaAttnModule, - Dsv4CsaAttnModule, Dsv4CsaTopkCalib, Dsv4MegamoeModule, MoeA2a, @@ -682,8 +680,6 @@ impl PerfDataFilename { Self::Dsv4CsaGenerationModule => "dsv4_csa_generation_module_perf.parquet", Self::Dsv4HcaGenerationModule => "dsv4_hca_generation_module_perf.parquet", Self::Dsv4PagedMqaLogitsModule => "dsv4_paged_mqa_logits_module_perf.parquet", - Self::Dsv4HcaAttnModule => "dsv4_hca_attn_module_perf.parquet", - Self::Dsv4CsaAttnModule => "dsv4_csa_attn_module_perf.parquet", Self::Dsv4CsaTopkCalib => "dsv4_csa_topk_calib_perf.parquet", Self::Dsv4MegamoeModule => "dsv4_megamoe_module_perf.parquet", Self::MoeA2a => "moe_a2a_perf.parquet", diff --git a/crates/core/src/perfmodel/operators/dsv4.rs b/crates/core/src/perfmodel/operators/dsv4.rs index 4e5ef60039..6a1eda6b83 100644 --- a/crates/core/src/perfmodel/operators/dsv4.rs +++ b/crates/core/src/perfmodel/operators/dsv4.rs @@ -1057,8 +1057,8 @@ mod tests { ); } - /// Real-DB fail-loud: b200_sxm/sglang/0.5.10 now retains only the - /// hca_attn table, so a CSA CP context query must fail LOUD with the + /// Real-DB fail-loud: b200_sxm/sglang/0.5.10 ships no DSV4 module tables + /// (its hca_attn sidecar was retired), so a CSA CP context query must fail LOUD with the /// actionable no-module-rows message (names the missing parquet) (never a fabricated value). /// The calib-specific missing-top_last branch is pinned synthetically in /// `perf_database::dsv4` (absent calib -> None -> operator fail-loud). diff --git a/crates/core/src/perfmodel/perf_database/dsv4.rs b/crates/core/src/perfmodel/perf_database/dsv4.rs index 9434a7a260..7bc151aa2b 100644 --- a/crates/core/src/perfmodel/perf_database/dsv4.rs +++ b/crates/core/src/perfmodel/perf_database/dsv4.rs @@ -537,9 +537,9 @@ impl Dsv4Table { /// table. Mirror of Python /// `ContextDeepSeekV4AttentionModule._lookup_sparse_kernel` with the /// kernel FIXED to `"paged_mqa_logits"` — the quadratic pair-count SOL - /// below is valid only for the causal indexer logits; the windowed - /// sidecars (hca_attn / csa_attn) must NOT route here (Python raises on - /// any other kernel name; Rust simply has no other entry point). + /// below is valid only for the causal indexer logits, the one sparse + /// kernel table left (the windowed hca_attn / csa_attn sidecars were + /// retired 2026-10-15; Rust has no other entry point). /// /// Table shape `[native_heads][tp][past_kv][isl][bs]`; `tp` falls back to /// 1 (the kernel is collected at tp=1 only and its work is diff --git a/crates/core/src/perfmodel/perf_database/table_view.rs b/crates/core/src/perfmodel/perf_database/table_view.rs index ffaa5f831d..78e04fa8bf 100644 --- a/crates/core/src/perfmodel/perf_database/table_view.rs +++ b/crates/core/src/perfmodel/perf_database/table_view.rs @@ -2448,9 +2448,9 @@ pub const TABLE_VIEW_ATTRIBUTES: &[(&str, &[&str])] = &[ /// The engine-side table view: `attribute` is the PerfDatabase attribute name /// the retired Python loader used to fill (`"_gemm_data"`, ...). Returns /// `Ok(None)` exactly when that loader returned `None` (every source path -/// missing), and the loader-shaped nested JSON otherwise. The three DSV4 -/// sparse sub-tables are addressed as -/// `"_dsv4_sparse_kernel_data."`. +/// missing), and the loader-shaped nested JSON otherwise. The DSV4 sparse +/// sub-table is addressed as `"_dsv4_sparse_kernel_data.paged_mqa_logits"` +/// (the hca_attn / csa_attn sidecars were retired 2026-10-15). pub fn table_view_json(tables: &PerfTables, attribute: &str) -> Result, AicError> { let src = |basename: &str| { tables diff --git a/docs/perf-model/collector/upgrade.md b/docs/perf-model/collector/upgrade.md index 1c3f2916e9..d5bd025ffa 100644 --- a/docs/perf-model/collector/upgrade.md +++ b/docs/perf-model/collector/upgrade.md @@ -32,17 +32,18 @@ Record these decisions before changing code: - what may be committed, pushed, or opened as a PR. `collector/framework_manifest.yaml` is the source of truth for the stock -framework pin and image. Module-level `__compat__` metadata may narrow that -contract but must not silently widen it. Treat stock and WideEP/special images +framework pin and image. Each collector module pins exactly the release the +manifest resolves for its op (the "Version contract" in +`python/aisimulate/.claude/rules/collector/layer_permissions.md`). Treat stock and WideEP/special images as separate runtimes with separate evidence. Do not silently broaden the work to another framework, XPU, WideEP, host upgrades, or performance-data publication. If the work is stacked on another PR, keep that dependency explicit when creating the branch and eventual PR. -For a fixed-release collector, prefer an exact version contract when the -framework API and kernel selection changed materially. Do not add compatibility -branches for versions that the project no longer intends to collect. +Upgrade collectors in place: move the exact `__compat__` pin to the new +release and do not keep compatibility branches for earlier releases; their +code stays in git history and their data stays under its version key. ## 2. Pin and prove the runtime @@ -163,7 +164,11 @@ on its observed speed. After the operation finishes, verify its checkpoint, summary, output shape, errors, container exit, and GPU state, then continue to the next operation without waiting for approval unless a real anomaly appears. -Use one output/checkpoint namespace per framework, version, platform, and +`tools/perf_database/collect_campaign.py` is the checked-in way to run this +stage: manifest-pinned image plus finalize deps, one namespace per declared +shard, `--resume` always and `--keep-csv` never (that flag skips finalization +and leaves no parquet or provenance), a GPU pool, and a `campaign.jsonl` work +log. Use one output/checkpoint namespace per framework, version, platform, and operation. Resume only when the current plan still contains the completed task IDs. If enumeration changes, prove which completed IDs remain valid and keep a hash-verified checkpoint backup before migration. diff --git a/python/aisimulate/.claude/rules/collector/failure_handling.md b/python/aisimulate/.claude/rules/collector/failure_handling.md index 45b84099d9..364d9a3a73 100644 --- a/python/aisimulate/.claude/rules/collector/failure_handling.md +++ b/python/aisimulate/.claude/rules/collector/failure_handling.md @@ -18,8 +18,10 @@ group is failing, it should be FIXED, not silently tolerated. ## What happens automatically (no action needed) - Every worker failure is recorded in `errors_.json` and - `collection_summary_.json` with `classification: "unexpected"`, - the case parameters, exception type/message, and a `(model, dtype)` group + `collection_summary_.json` with a classification: `kernel_limit` + when the message starts with `FIXME(kernel-limit)` (a cited guard, see + layer_permissions.md), otherwise `unexpected`. Each record carries the + case parameters, exception type/message, and a `(model, dtype)` group label. The summary aggregates `errors_by_group`; a group with many failures is logged as a fix-me warning. - CUDA-fatal errors reset the worker process; the task is already recorded diff --git a/python/aisimulate/.claude/rules/collector/layer_permissions.md b/python/aisimulate/.claude/rules/collector/layer_permissions.md index 72c2722aed..d20caff886 100644 --- a/python/aisimulate/.claude/rules/collector/layer_permissions.md +++ b/python/aisimulate/.claude/rules/collector/layer_permissions.md @@ -150,6 +150,45 @@ prediction) or a guard that raises a classified error citing the framework source line, or deletes the note. Do not implement guards from unverified claims, and never express these limits in YAML. +**Evidence-backed guards** (owner decision 2026-10-30). Some limits cannot be +shown in source (a precompiled cubin, a JIT kernel generator inside a shared +library), and some failures are expensive (a worker-killing abort, a full +module load per case). A guard for such a limit is legal only when it: + +1. **Raises at run time, in the worker, before the expensive step.** Never at + plan time: the case stays queued and lands in the failure log as + `kernel_limit`, so every run still shows it. +2. **Starts its message with `FIXME(kernel-limit):`.** +3. **Cites its evidence.** The framework source line when one exists; + otherwise a reproducible failing run (job or shard, collector revision, + failure counts) whose passing rows bound the guarded envelope exactly. A + single failing case or an unbounded envelope is not evidence. +4. **Is re-verified on every version bump** like any other + `FIXME(kernel-limit)`, by a source check or by one run without it. + +## Version contract: one release per collector file + +Owner policy (2026-09-20, made binding 2026-10-30): version history lives in +git, not in side-by-side files or version ranges. + +1. **Each collector module targets exactly one framework release.** Its + `__compat__` is a single exact pin, `"=="`, equal to + the runtime `framework_manifest.yaml` resolves for its op (the framework + default or a family pin). No ranges, no `!=` exclusions, no per-version + branches for releases the manifest does not collect. Per-architecture + gates inside a module may only narrow that one release. +2. **Upgrade in place.** A version bump rewrites the module for the new + release and moves its pin. Prior-version code is `git log -- ` + away; prior-version data stays permanent under its version key. Do not + add `_` fork modules or keep old releases admitted "so they + still work". +3. **Pin consolidation.** A model or family pin below the target version + upgrades with it; a pin above the default is the next version for that + family and needs its own GPU validation. +4. **Enforced by tests.** `tests/unit/collector/test_framework_manifest.py` + fails when a stock vLLM, SGLang or TRT-LLM registry module declares + anything but `==` its resolved runtime. + ## Module boundary: a collector task may only touch collector/ The collector's outward surface is its **data contract**: the perf row schema @@ -186,20 +225,27 @@ A rule in this file may be temporarily excepted ONLY through an entry in the only when the entry carries ALL of: 1. **A self-contained machine-readable scope** — the rule excepted, plus the - exact tables/backend/versions/systems the exception covers. Prose-only - scoping is invalid. + exact tables/backend/systems the exception covers. Prose-only scoping is + invalid. `versions: all` (the default; owner decision 2026-10-30) covers + every version of the backend: an exception follows the deviation, not a + release, and re-verification (item 3) is what bounds it. A version list + narrows it. 2. **The standards owner's `approved_by` sign-off.** These rule files are human-owned policy; only their owner can except them. -3. **An `expires` date and a linked retirement work item.** An exception is - a debt with a due date, not a policy change; on expiry the exception is - renewed by the owner or the deviation is fixed. +3. **A `verified_on` release and a linked retirement work item** (owner + decision 2026-10-30; replaces the former `expires` date). An exception + is a debt, not a policy change. Every bump of the collector pin for the + covered tables re-checks the deviation against the new release, by a + source check or a run, as for a `FIXME(kernel-limit)`: if it still + holds, `verified_on` moves to that release; if not, the entry is + deleted. An entry whose `verified_on` is behind the pin is stale. 4. **Consumer-visible provenance** — the deviation must be discoverable from the shipped data's own declaration layer (e.g. the `reuse.yaml` reason or sidecar for the affected tables), never only from this registry. An exception documents and bounds a deviation; it does not weaken the rule. -Reviews treat an in-scope, unexpired entry as compliant and everything -outside its scope as an ordinary violation. +Reviews treat an in-scope entry verified on the current pin as compliant; +a stale entry, and everything outside its scope, is an ordinary violation. ## Meta rules diff --git a/python/aisimulate/.claude/skills/aic-auto-collect/SKILL.md b/python/aisimulate/.claude/skills/aic-auto-collect/SKILL.md index 23ff177e6c..c630866cf2 100644 --- a/python/aisimulate/.claude/skills/aic-auto-collect/SKILL.md +++ b/python/aisimulate/.claude/skills/aic-auto-collect/SKILL.md @@ -239,6 +239,37 @@ python3 collector/collect.py --backend --ops \ --checkpoint-dir "$COLLECTOR_CHECKPOINT_DIR///" --resume ``` +### Campaign mode (full plan, several GPUs) — use the checked-in launcher, do not hand-roll one + +Once an op's smoke/limited runs are clean, run the full retained plan through +`tools/perf_database/collect_campaign.py`. It fixes the parameters that every +hand-written launcher has drifted on (sm120 campaign 2026-09-30: `--keep-csv` +skipped finalize for 13 shards; the official image lacked pyarrow): + +```bash +# 1. image = the framework_manifest pin for (backend, op) + the finalize deps (pyarrow, pandas) +python3 tools/perf_database/collect_campaign.py build-image --backend --op [--build-arg http_proxy=...] +# 2. declare the shards once (YAML: backend, sm, shards[{op, name, case_filter}]) — one shard = one +# output/checkpoint namespace ///; case_filter is collect.py's substring filter +# 3. run the GPU pool; re-run the same command after any interruption (DONE shards skip, others --resume) +python3 tools/perf_database/collect_campaign.py run --plan shards.yaml --root /collect --gpus 0,1,2,3 +python3 tools/perf_database/collect_campaign.py status --plan shards.yaml --root /collect +# 4. shards that were produced by hand with --keep-csv have no parquet: finalize them (resume-only re-entry) +python3 tools/perf_database/collect_campaign.py finalize --plan shards.yaml --root /collect --gpus 0 +``` + +Parameter rules the tool enforces (and that apply to any manual `collect.py` run too): + +- `--resume` always, `--keep-csv` never for data you intend to deliver: `--keep-csv` means + "keep CSV *instead of* finalizing" — no parquet, no `collection_meta.yaml`, no provenance. + Its only use is inspecting raw staging rows. +- `--sm ` explicit on every run (never rely on auto-detect inside a container). +- Never point two concurrent runs at one cwd: staging files are named per table + (`gemm_perf.txt`) and would interleave. Shards of one op are merged at delivery + (playbook §8), not by sharing a directory. +- The image is the manifest pin for the op (family pins can differ per op) — a + hand-typed tag is a review finding. + Use a separate log directory per op when iterating: ```bash @@ -305,17 +336,16 @@ Fix policy: 2. Read the failing `collector//collect_*.py`. 3. Search the framework source checkout for the missing class, function, import path, or attribute. 4. Verify the real API in source. Do not trust Python's "Did you mean" suggestions blindly. - 5. Patch the collector with a minimal compatibility-preserving change. -- For API changes, use version routing and `__compat__` where needed. + 5. Patch the collector in place for the target release (one release per file: `.claude/rules/collector/layer_permissions.md`, "Version contract"). +- For API changes, upgrade the collector in place and pin `__compat__` to exactly the target release. - For GPU capability gaps, filter test cases before execution using `get_sm_version()`. - For Blackwell/SM100+ features, gate FP4/NVFP4/FP8 paths carefully. - For dimension constraints, filter on per-rank dimensions such as `inter_size // tp`. - For CUDA-fatal errors, prefer worker restart and skip only the exact invalid region. - Preserve data quality. Do not reduce benchmark repetitions or broad test dimensions just to pass. -- Keep older backend versions working. - Keep shared test-case generators deterministic when unit tests depend on them. Put runtime availability probes in backend collector wrappers when the skip is caused by a partially installed framework package. - Guard `importlib.util.find_spec("a.b.c")` with `try/except ModuleNotFoundError`; partial packages can make `find_spec` raise instead of return `None`. -- If a collector route only exists for future framework versions, use `VersionRoute` in `collector//registry.py` and add a wrapper module with an explicit `__compat__`. +- Do not add side-by-side version fork modules; prior-version code is `git log -- ` away. - If a collector helper silently returns `None` after a dry-run or subprocess failure, make it raise or record the skip explicitly so failed coverage is visible. - If a valid production-like case fails in the framework kernel, file or prepare an upstream bug for SGLang, vLLM, or TRT-LLM with exact image/version, SM, minimal repro, and traceback. Link it in the PR or failure ledger. - If the case is not production-like or the runtime explicitly does not support that SM/path, skip it in `collect_xx.py` before execution with a clear comment and a narrow predicate. diff --git a/python/aisimulate/.claude/skills/aic-collector-op-development/SKILL.md b/python/aisimulate/.claude/skills/aic-collector-op-development/SKILL.md index 5cf9579748..3df1b79ae2 100644 --- a/python/aisimulate/.claude/skills/aic-collector-op-development/SKILL.md +++ b/python/aisimulate/.claude/skills/aic-collector-op-development/SKILL.md @@ -54,8 +54,9 @@ If the task proceeds to GPU collection, also use `$aic-auto-collect`. plan selection required by the resolved case plan. Generic resume, retry, checkpoint, logging, and output-finalization behavior require separate explicit scope. -- Anchor collector behavior to the requested framework version. Prefer one - verified `__compat__`/version route over speculative multi-version branches. +- Anchor collector behavior to the requested framework version: pin exactly + one release (`__compat__ = "=="`, see the "Version contract" + section of `.claude/rules/collector/layer_permissions.md`). ## Step 1: Establish the consumer contract diff --git a/python/aisimulate/collector/README.md b/python/aisimulate/collector/README.md index 96834013c9..a7e8d0ffae 100644 --- a/python/aisimulate/collector/README.md +++ b/python/aisimulate/collector/README.md @@ -451,15 +451,28 @@ If unavailable, a warning is logged and execution continues without power data. - Kernel iterations are automatically adjusted to meet minimum duration for accurate measurements - Backward compatible: without `--measure_power`, CSVs remain unchanged -## CUDA Graph Fallback Support - -The `benchmark_with_power` helper function now supports graceful fallback to eager execution when CUDA graph capture fails. This is particularly useful for complex operations like MOE (Mixture of Experts) with large batch sizes. - -### Features -- **Automatic fallback**: When `allow_graph_fail=True`, CUDA graph capture failures trigger eager execution instead of raising exceptions -- **Power measurement in both paths**: Power monitoring works correctly in both graph replay and eager execution modes -- **Memory safety**: Automatic `torch.cuda.empty_cache()` call on graph capture failure to prevent memory fragmentation -- **Transparency**: Results include `used_cuda_graph` flag to indicate which execution path was used +## CUDA Graph Timing Policy + +`benchmark_with_power` times kernels under CUDA-graph replay. A failed capture +raises (`allow_graph_fail=False`, the default) and the case fails classified. + +- **No collector that writes a perf table may pass `allow_graph_fail=True`.** + The eager fallback times the Python launch loop and ships that floor as the + kernel latency, and no perf table records the helper's `used_cuda_graph` + flag — `.claude/rules/collector/layer_permissions.md` allows a + measurement-method degradation only when it is recorded in the output row. + `tests/unit/collector/test_no_eager_graph_fallback_sites.py` pins this for + every file under `collector/`. The last fallback sites (trtllm mamba2 and + moe, wideep moe_compute, network trtllm alltoall) were closed on 2026-10-12; + the sglang trtllm_mla prefill (B1) and the trtllm GDN lanes the day before. +- A collector that sweeps several cells inside one task counts failed cells + and raises a strict-completeness error at the end (GDN, mamba2), so the + other cells still ship rows and the failure lands in `errors_.json`. +- Multi-rank collectors agree on the capture outcome before any rank replays + (`collect_trtllm_alltoall._agree_graph_capturable`); a one-sided capture + failure would otherwise hang the peers' first replay. +- `allow_graph_fail=True` remains available for ad-hoc diagnosis only; the + result then carries `used_cuda_graph=False` and must not be persisted. ### Usage Example ```python @@ -469,33 +482,17 @@ def my_kernel(): # Your kernel code here moe.forward(hidden_states, logits) -# Use benchmark_with_power with fallback support with benchmark_with_power( device=device, kernel_func=my_kernel, num_warmups=3, num_runs=6, repeat_n=1, - allow_graph_fail=True, # Enable graceful fallback ) as results: latency = results["latency_ms"] - power_stats = results["power_stats"] # Available in both paths - - # Check which execution path was used - if not results["used_cuda_graph"]: - print("CUDA graph capture failed, used eager execution") + power_stats = results["power_stats"] ``` -### When to Use -- **Complex operations**: MOE, dynamic memory patterns, or operations that may not be graph-compatible -- **Large batch sizes**: When graph capture may fail due to memory constraints -- **Development/debugging**: To ensure collection continues even if graph capture fails - -### Backward Compatibility -- Default behavior unchanged: `allow_graph_fail=False` maintains existing behavior -- Existing collectors work without modifications -- Only opt-in when needed for specific use cases - ## For TensorRT-LLM ### Optional( Read if collecting for mxfp4 kernels) To collect performance data for mxfp4 kernels used in GPTOSS models, depending on the version of TensorRT-LLM, you might need to manually install triton-kernels. If your version is **>= 1.3.0rc2**, nothing needs to be done and you can run the collection process directly. For any version before **1.3.0rc2**. Follow the below instructions to install triton-kernels and expose them to TensorRT-LLM diff --git a/python/aisimulate/collector/case_generator.py b/python/aisimulate/collector/case_generator.py index e030be2ab9..ce931dd199 100644 --- a/python/aisimulate/collector/case_generator.py +++ b/python/aisimulate/collector/case_generator.py @@ -715,6 +715,12 @@ class MLAModuleModelSpec: architecture: str native_num_heads: int wideep_mla: bool + # Declared real tensor-parallel shards of this model (case_authoring's + # model-correlated narrowing). When set, the MSA collectors derive their + # per-rank head axis as native_num_heads // tp instead of sweeping the + # shared inner_sweep_head_counts grid (owner decision 2026-09-29: MSA is + # one fixed geometry, MiniMax-M3 64q/4kv/4idx; TP 16 is the ceiling). + tensor_parallel_sizes: tuple[int, ...] | None = None @dataclasses.dataclass(frozen=True) @@ -745,6 +751,19 @@ class MLAModuleSweepSpec: # memory contract (their guard tests pin it). None -> fall back to # generation_max_tokens. generation_msa_max_tokens: int | None = None + # Optional MSA context budget for the TRT-LLM Triton reference + # implementation (cases/base_ops/mla_module.yaml mla_module_trtllm.context + # .msa_triton_max_tokens): the collector applies it only where it resolves + # implementation="triton" (off SM100/103); fmha_sm100 systems keep + # context_max_tokens. None -> no extra cap. + context_msa_triton_max_tokens: int | None = None + # Optional DSA-specific context budget: DSA context follows serving's + # chunked-prefill shape (new-token budget = chunk size, long context via + # prefix). None on all three -> the DSA branch falls back to the shared + # context axes (context_max_tokens / context_prefix_lengths). + context_dsa_chunk_prefill_size: int | None = None + context_dsa_max_full_sequence_length: int | None = None + context_dsa_prefix_lengths: list[int] | None = None @dataclasses.dataclass(frozen=True) @@ -790,13 +809,28 @@ def get_mla_module_model_specs( continue if wideep_mla is not None and bool(value.get("wideep_mla", False)) != wideep_mla: continue + native_num_heads = int(value["native_num_heads"]) + tp_sizes = None + if value.get("tensor_parallel_sizes") is not None: + tp_list = _as_int_list( + value.get("tensor_parallel_sizes"), + field_name="model_case_values.mla_module.tensor_parallel_sizes", + ) + for tp in tp_list: + if tp <= 0 or native_num_heads % tp != 0: + raise ValueError( + "model_case_values.mla_module.tensor_parallel_sizes must be positive divisors of " + f"native_num_heads: model={value['model_path']}, native_num_heads={native_num_heads}, tp={tp}" + ) + tp_sizes = tuple(tp_list) specs.append( MLAModuleModelSpec( model_path=str(value["model_path"]), attention_type=str(value["attention_type"]), architecture=str(value["architecture"]), - native_num_heads=int(value["native_num_heads"]), + native_num_heads=native_num_heads, wideep_mla=bool(value.get("wideep_mla", False)), + tensor_parallel_sizes=tp_sizes, ) ) @@ -847,6 +881,28 @@ def get_mla_module_model_specs( return specs +def get_msa_head_counts(backend: str | None = None) -> list[int]: + """Per-rank head axis for the MSA module collectors, descending. + + Derived from the declared MSA model rows: ``native_num_heads // tp`` for + each ``tensor_parallel_sizes`` entry (falling back to the sweep's + ``module_tp_sizes`` when a row declares none). MSA serves one fixed + geometry, so unlike GQA attention there is no shared head grid to sweep; + head counts no real shard produces (48/96 -> 3/6 index heads) are never + generated. + """ + sweep = get_mla_module_sweep_spec(backend) + heads: set[int] = set() + for spec in get_mla_module_model_specs(attention_type="msa", backend=backend): + tp_sizes = spec.tensor_parallel_sizes or tuple(sweep.module_tp_sizes) + for tp in tp_sizes: + if spec.native_num_heads % tp == 0: + heads.add(spec.native_num_heads // tp) + if not heads: + raise ValueError("mla_module: no MSA model row yields a per-rank head count (check tensor_parallel_sizes)") + return sorted(heads, reverse=True) + + def _required_mapping(value: object, *, field_name: str) -> dict[str, object]: if not isinstance(value, dict): raise TypeError(f"{field_name} must be a mapping") @@ -1013,6 +1069,7 @@ def get_mla_module_sweep_spec(backend: str | None = None) -> MLAModuleSweepSpec: context_max_tokens=int(context["max_tokens"]), context_large_sequence_min=_optional_int(context.get("large_sequence_min")), context_large_sequence_max_batch_size=_optional_int(context.get("large_sequence_max_batch_size")), + context_msa_triton_max_tokens=_optional_int(context.get("msa_triton_max_tokens")), generation_max_tokens=int(generation["max_tokens"]), generation_msa_max_tokens=_optional_int(generation.get("msa_max_tokens")), generation_large_sequence_min=_optional_int(generation.get("large_sequence_min")), @@ -1023,6 +1080,16 @@ def get_mla_module_sweep_spec(backend: str | None = None) -> MLAModuleSweepSpec: field_name="mla_module.context.prefix_lengths", default=[0], ), + context_dsa_chunk_prefill_size=_optional_int(context.get("dsa_chunk_prefill_size")), + context_dsa_max_full_sequence_length=_optional_int(context.get("dsa_max_full_sequence_length")), + context_dsa_prefix_lengths=( + _optional_int_list( + context.get("dsa_prefix_lengths"), + field_name="mla_module.context.dsa_prefix_lengths", + default=[], + ) + or None + ), ) @@ -2475,10 +2542,6 @@ def _dsv4_module_budgets() -> dict: _DSV4_SPARSE_CHUNK_PREFILL_SIZE = int(_DSV4_CONFIG["sparse_chunk_prefill_size"]) _DSV4_SPARSE_MAX_FULL_S = int(_DSV4_CONFIG["sparse_max_full_sequence_length"]) _DSV4_SPARSE_TP_SIZES = _required_mapping(_DSV4_CONFIG["sparse_tp_sizes"], field_name="dsv4.sparse_tp_sizes") -_DSV4_SPARSE_TP_LIST_ATTN = _as_int_list( - _DSV4_SPARSE_TP_SIZES["hca_attn"], - field_name="dsv4.sparse_tp_sizes.hca_attn", -) _DSV4_SPARSE_TP_LIST_INDEXER = _as_int_list( _DSV4_SPARSE_TP_SIZES["paged_mqa_logits"], field_name="dsv4.sparse_tp_sizes.paged_mqa_logits", @@ -2695,8 +2758,8 @@ def get_dsv4_hca_generation_test_cases(): def get_dsv4_topk_calib_test_cases(): """One case per model (the topk DELTA calibration is just another member of - the sparse-op family — same ``[model_path, kernel]`` shape as paged_mqa / - csa_attn / hca_attn, dispatched through ``run_dsv4_sparse_kernel_worker``). + the sparse-op family — same ``[model_path, kernel]`` shape as + paged_mqa_logits, dispatched through ``run_dsv4_sparse_kernel_worker``). The grid matches the CSA module data 1:1: the worker reads the already-collected ``dsv4_csa_*_module_perf.txt`` and benches exactly those @@ -2709,7 +2772,11 @@ def get_dsv4_topk_calib_test_cases(): return [[model_path, "topk"] for model_path in _selected_dsv4_calib_models()] -DSV4_SPARSE_KERNELS = ("paged_mqa_logits", "hca_attn") +# The CSA indexer scoring GEMM is the only sparse sub-kernel collected standalone. +# hca_attn / csa_attn (FMLA sidecars) were removed 2026-10-15: retired SDK +# placeholders with no query path, and registry-only ops that failed every +# shard plan generated from the registry. +DSV4_SPARSE_KERNELS = ("paged_mqa_logits",) def _build_dsv4_sparse_test_cases( @@ -2717,7 +2784,6 @@ def _build_dsv4_sparse_test_cases( bs_list=None, isl_list=None, past_kv_list=None, - tp_list_attn=None, tp_list_indexer=None, ): """Generate ``(bs, isl, past_kv, tp_size, kernel, model)`` tuples. @@ -2729,14 +2795,12 @@ def _build_dsv4_sparse_test_cases( bs_list = list(bs_list) if bs_list is not None else list(_DSV4_SPARSE_BS_LIST) isl_list = list(isl_list) if isl_list is not None else list(_DSV4_SPARSE_ISL_LIST) past_kv_list = list(past_kv_list) if past_kv_list is not None else list(_DSV4_SPARSE_PAST_KV_LIST) - tp_list_attn = list(tp_list_attn) if tp_list_attn is not None else list(_DSV4_SPARSE_TP_LIST_ATTN) tp_list_indexer = list(tp_list_indexer) if tp_list_indexer is not None else list(_DSV4_SPARSE_TP_LIST_INDEXER) cases = [] for model_path in _selected_dsv4_models(): for kernel in kernels: - tp_list = tp_list_attn if kernel == "hca_attn" else tp_list_indexer - for tp_size in tp_list: + for tp_size in tp_list_indexer: for bs in bs_list: for isl in isl_list: if bs * isl > _DSV4_SPARSE_CHUNK_PREFILL_SIZE: @@ -2747,8 +2811,6 @@ def _build_dsv4_sparse_test_cases( full_s = isl + past_kv if kernel == "paged_mqa_logits" and full_s < 4: continue - if kernel == "hca_attn" and full_s < 64: - continue cases.append( [ bs, @@ -2769,19 +2831,6 @@ def get_dsv4_paged_mqa_logits_test_cases(): return [[model_path, "paged_mqa_logits"] for model_path in _selected_dsv4_models()] -def get_dsv4_hca_attn_test_cases(): - """One case per model. The worker derives shapes from the HCA module CSV - (hca_attn is the HCA c128 FMLA sub-kernel), 1:1 with the hca-module rows.""" - return [[model_path, "hca_attn"] for model_path in _selected_dsv4_models()] - - -def get_dsv4_csa_attn_test_cases(): - """One case per model. The worker derives shapes from the CSA module CSV - (csa_attn is the CSA sparse FMLA over the topk-selected c4 positions), - 1:1 with the csa-module rows — same source as paged_mqa_logits/topk.""" - return [[model_path, "csa_attn"] for model_path in _selected_dsv4_models()] - - # Backward-compatible names for older PR comments/tests while the registry and # persisted perf files use the upstream DeepSeek-V4 op names. _DSV4_FLASH_CONFIG = _DSV4_CONFIG @@ -2796,7 +2845,6 @@ def get_dsv4_csa_attn_test_cases(): _DSV4_FLASH_SPARSE_PAST_KV_LIST = _DSV4_SPARSE_PAST_KV_LIST _DSV4_FLASH_SPARSE_CHUNK_PREFILL_SIZE = _DSV4_SPARSE_CHUNK_PREFILL_SIZE _DSV4_FLASH_SPARSE_MAX_FULL_S = _DSV4_SPARSE_MAX_FULL_S -_DSV4_FLASH_SPARSE_TP_LIST_ATTN = _DSV4_SPARSE_TP_LIST_ATTN _DSV4_FLASH_SPARSE_TP_LIST_INDEXER = _DSV4_SPARSE_TP_LIST_INDEXER DSV4_FLASH_SPARSE_KERNELS = DSV4_SPARSE_KERNELS _dsv4_flash_config = _dsv4_config @@ -2819,4 +2867,3 @@ def _dsv4_flash_model_path() -> str: get_dsv4_flash_csa_generation_test_cases = get_dsv4_csa_generation_test_cases get_dsv4_flash_hca_generation_test_cases = get_dsv4_hca_generation_test_cases get_dsv4_flash_paged_mqa_logits_test_cases = get_dsv4_paged_mqa_logits_test_cases -get_dsv4_flash_hca_attn_test_cases = get_dsv4_hca_attn_test_cases diff --git a/python/aisimulate/collector/cases/base_ops/mla_module.yaml b/python/aisimulate/collector/cases/base_ops/mla_module.yaml index b17830b580..a0540d2240 100644 --- a/python/aisimulate/collector/cases/base_ops/mla_module.yaml +++ b/python/aisimulate/collector/cases/base_ops/mla_module.yaml @@ -266,6 +266,23 @@ common_case_values: - 1 context: max_tokens: 131072 + # MSA (MiniMax-M3) on TRT-LLM has two implementations: the fmha_sm100 + # kernels (implementation="msa", SM100/103 only — msa_availability.py + # ensure_msa_available@1.3.0rc29 raises everywhere else, SM120 included) + # and the Triton reference (implementation="triton", the only path on + # SM89/SM90/SM120). The reference computes the block-sparse GQA as a + # masked DENSE FP32 attention with per-query K/V expansion + # (triton_backend.py:_sparse_gqa_masked@1.3.0rc29), so its context cost + # is quadratic per sequence: 64 heads, isl 32768, b=1 takes 24.5 s per + # forward on rtxpro6000 and 18 s on h200_sxm (fmha_sm100 on gb200: 9 ms), + # and b*isl=131072 OOMs on the unchunked FP32 index-score slab (64 GiB). + # Owner decision 2026-10-11: on the Triton reference path cap context at + # batch*isl <= 32768 (isl 32768 itself stays, at b=1); the fmha_sm100 + # systems keep the shared 131072. Applied by collect_msa_module's + # context getter where the collector resolves implementation="triton"; + # re-verify on every TRT-LLM version bump (a real sparse prefill kernel + # for these SMs lifts it). + msa_triton_max_tokens: 32768 prefix_lengths: - 0 - 128 @@ -378,20 +395,10 @@ common_case_values: min_sm: 89 phases: [context, generation] attention_types: [mla, dsa] - # Native tensor-scaled FP8 weights with dynamic activations; distinct - # from block-scaled FP8 and from the fp8_static GEMM contract. - - compute_dtype: bfloat16 - kv_cache_dtype: bfloat16 - gemm_type: fp8 - min_sm: 89 - phases: [context, generation] - attention_types: [mla, dsa] - - compute_dtype: bfloat16 - kv_cache_dtype: fp8 - gemm_type: fp8 - min_sm: 90 - phases: [context, generation] - attention_types: [mla, dsa] + # No per-tensor fp8 GEMM combo (removed 2026-10-30): every bundled MLA/DSA + # FP8 checkpoint is block-scaled (fp8_block) and the only consumer was the + # prediction gate's forced per-tensor override. Existing tables keep their + # rows. Same reasoning as the AIC-1782 Task V2 MoE fp8 removal. - compute_dtype: bfloat16 kv_cache_dtype: fp8 gemm_type: fp8_block @@ -464,6 +471,47 @@ common_case_values: prefix_lengths: - 0 - 128 + # DSA context follows serving's chunked-prefill shape, matching the + # DeepSeek-V4 sparse + sglang-DSA convention (DeepseekV4ForCausalLM_cases + # sparse_chunk_prefill_size / sparse_past_kv_lengths; mla_module_sglang + # context.prefix_lengths). The per-forward NEW-token budget is the chunk + # size and the long context enters through prefix, so (a) each row is + # faithful to what the engine actually queries — context_attention at + # isl=chunk, prefix=context, engine/runtime.rs mixed_step Pass 2, NOT a + # monolithic isl forward serving never runs — and (b) the fp8 sparse + # decode kernel's q stride (bs*isl*padded_heads*576, mixed mode at + # num_heads<32) stays under the int32 limit (overflow at bs*isl>~58k; + # the 776 dsa_context fp8 failures, 2026-09-22). MLA/MSA are unaffected: + # only the collector's DSA branch sweeps prefix. + dsa_chunk_prefill_size: 16384 + dsa_max_full_sequence_length: 1048576 + dsa_prefix_lengths: + - 0 + - 1 + - 4 + - 8 + - 16 + - 32 + - 64 + - 128 + - 256 + - 512 + - 1024 + - 1536 + - 2048 + - 3072 + - 4096 + - 6144 + - 8192 + - 10240 + - 12288 + - 16384 + - 32768 + - 65536 + - 131072 + - 262144 + - 524288 + - 1048575 generation: max_tokens: 33554432 large_cache_tokens: 16777216 diff --git a/python/aisimulate/collector/cases/capabilities.yaml b/python/aisimulate/collector/cases/capabilities.yaml index 68b0589811..677fb7dc83 100644 --- a/python/aisimulate/collector/cases/capabilities.yaml +++ b/python/aisimulate/collector/cases/capabilities.yaml @@ -47,11 +47,19 @@ op_min_sm: # The DSV4/GLM-5 sparse sub-kernel collectors depend on the same SM90+ # kernel libraries (DeepGEMM fp8_mqa_logits, FlashMLA sparse, jit topk). dsv4_paged_mqa_logits_module: 90 - dsv4_hca_attn_module: 90 - dsv4_csa_attn_module: 90 dsv4_csa_topk_calib: 90 glm5_mqa_logits_module: 90 glm5_topk_module: 90 glm5_dsa_attn_module: 90 + # The DSV4 attention module collectors run the production DSV4 attention class, whose + # only backends are FlashMLA (SM90 / SM10x) and the SM100+ paths: on SM89 the selector + # falls back to a FlashMLA that does not support the arch, and serving fails the same way + # (vLLM 0.30.0 probe on L40, DeepGEMM hyperconnection "Unsupported architecture", + # collector/opharness findings sm89_*). Before this entry every sm89 case was queued and + # failed one by one as an `unexpected` RuntimeError. + dsv4_csa_context_module: 90 + dsv4_csa_generation_module: 90 + dsv4_hca_context_module: 90 + dsv4_hca_generation_module: 90 # compute_scale measures FP8 quantization overhead. compute_scale: 89 diff --git a/python/aisimulate/collector/cases/models/DeepseekV32ForCausalLM_cases.yaml b/python/aisimulate/collector/cases/models/DeepseekV32ForCausalLM_cases.yaml index f7e3b461c7..dedd779d84 100644 --- a/python/aisimulate/collector/cases/models/DeepseekV32ForCausalLM_cases.yaml +++ b/python/aisimulate/collector/cases/models/DeepseekV32ForCausalLM_cases.yaml @@ -20,6 +20,11 @@ model_case_values: - model_path: deepseek-ai/DeepSeek-V3.2 attention_type: dsa native_num_heads: 128 + # Real tensor-parallel shards (case_authoring model-correlated narrowing): + # module collectors that honour it sweep per-rank heads = 128 // tp + # (128/64/32/16/8) instead of the shared inner head grid, whose + # 96/48/24/12 entries belong to the 96-head Kimi-K3 MLA module. + tensor_parallel_sizes: [1, 2, 4, 8, 16] all_frameworks_op_cases: moe: @@ -31,8 +36,6 @@ framework_specific_op_cases: cases: all dsa_generation_module: cases: all - moe_ep: - cases: all vllm: dsa_context_module: cases: all diff --git a/python/aisimulate/collector/cases/models/DeepseekV3ForCausalLM_cases.yaml b/python/aisimulate/collector/cases/models/DeepseekV3ForCausalLM_cases.yaml index 515b85edbd..d1671115a1 100644 --- a/python/aisimulate/collector/cases/models/DeepseekV3ForCausalLM_cases.yaml +++ b/python/aisimulate/collector/cases/models/DeepseekV3ForCausalLM_cases.yaml @@ -13,9 +13,10 @@ base_ops: - gemm description: > DeepSeek-V3-family model-centric collection plan. Includes shared base op cases - plus stock MoE and MLA cases. SGLang WideEP ops (moe_ep) require a separate - runtime and must be requested independently; trtllm moe_ep pins the same image - as stock trtllm and is activated below. + plus stock MoE and MLA cases. WideEP ops (moe_ep) live on their own runtime + pins (wideep_sglang 0.5.10, wideep_trtllm 1.3.0rc20 vs stock trtllm 1.3.0rc29) + and are requested independently of this plan on both backends (2026-10-04: + the trtllm same-pin assumption was dropped when stock trtllm moved to rc29). model_case_values: moe: @@ -108,8 +109,6 @@ framework_specific_op_cases: cases: all mla_generation_module: cases: all - moe_ep: - cases: all vllm: mla_context_module: cases: all diff --git a/python/aisimulate/collector/cases/models/DeepseekV4ForCausalLM_cases.yaml b/python/aisimulate/collector/cases/models/DeepseekV4ForCausalLM_cases.yaml index 56c2e2ff86..f659854c6a 100644 --- a/python/aisimulate/collector/cases/models/DeepseekV4ForCausalLM_cases.yaml +++ b/python/aisimulate/collector/cases/models/DeepseekV4ForCausalLM_cases.yaml @@ -393,8 +393,6 @@ model_case_values: sparse_chunk_prefill_size: 8192 sparse_max_full_sequence_length: 1048576 sparse_tp_sizes: - hca_attn: - - 1 paged_mqa_logits: - 1 @@ -419,17 +417,31 @@ framework_specific_op_cases: # after the csa module entries). dsv4_csa_topk_calib: cases: all - # Standalone paged_mqa/hca_attn/csa_attn rows are experimental-only: the - # production SDK has no query call-site for them. Keep their registry - # entries available for explicit research runs, but do not schedule them in - # the default model plan. + # paged_mqa_logits IS consumed: it is the only sparse-kernel sidecar with a + # live SDK query path (sdk/operations/dsv4.py: the CP chunk walk loads + # PerfDataFilename.dsv4_paged_mqa_logits_module) and every shipped platform + # carries the table. The 2026-10-05 b200 shard run lost it because this + # file still called it experimental (job 469988017). The hca_attn / + # csa_attn sidecars were removed from the collector and SDK on 2026-10-15: + # their SDK readers had been retired (csa_attn never shipped data) and the + # registry-only ops failed every sm100 shard plan generated from the + # registry since 2026-10-05. + dsv4_paged_mqa_logits_module: + cases: all mhc_module: cases: all vllm: # These four full-module paths are consumed by the production DSV4 model. - # Standalone sparse kernels and mHC remain registry-only because their - # current consumer contract/schedule does not represent vLLM 0.24.0's - # fused TileLang post->pre+norm and boundary/head calls. + # hca_attn / csa_attn were removed 2026-10-15 (see the sglang note above); + # paged_mqa_logits is consumed (see the sglang note above) and shipped for + # vllm 0.24/0.25. mHC is scheduled again (2026-10-11): the SDK queries + # mhc pre/post rows on every backend, collector/vllm/collect_mhc_module.py + # measures vLLM's own mhc_pre_tilelang / mhc_post_tilelang (qualified on + # B200 0.25, shipped 0.22/0.24/0.25) and documents why pre(no-norm)+post + # tracks serving's fused post+pre+norm boundary within 1-2%; the earlier + # "contract does not represent the fused TileLang path" note predated that + # measurement. Without this row vLLM 0.30 priced mHC from the 0.25 table + # through the nearest-earlier fallback. dsv4_csa_context_module: cases: all dsv4_hca_context_module: @@ -438,6 +450,10 @@ framework_specific_op_cases: cases: all dsv4_hca_generation_module: cases: all + dsv4_paged_mqa_logits_module: + cases: all + mhc_module: + cases: all trtllm: # TRT-LLM 1.3.0rc20 ships the mHC module + torch.ops.trtllm.mhc_* kernels # (tensorrt_llm/_torch/modules/mhc/) but no DeepseekV4 model class yet, so diff --git a/python/aisimulate/collector/cases/models/Gemma4ForConditionalGeneration_cases.yaml b/python/aisimulate/collector/cases/models/Gemma4ForConditionalGeneration_cases.yaml index 8af12e525e..7c6312c19d 100644 --- a/python/aisimulate/collector/cases/models/Gemma4ForConditionalGeneration_cases.yaml +++ b/python/aisimulate/collector/cases/models/Gemma4ForConditionalGeneration_cases.yaml @@ -25,6 +25,10 @@ description: > model_case_values: attention: + # sglang_backends: sglang 0.5.21 arg_groups/model_overrides/gemma4.py:_gemma4_overrides forces + # triton on every SM except SM100 (trtllm_mha); a missing SM key falls through to the generic + # table (flashinfer on sm89/sm120), which serving never runs for Gemma4. sm89 added 2026-10-04 + # (L40: gate attn_ctx_hd512_gemma-4-26B-A4B diverged — collector FlashInfer prefill vs serving Triton). # trtllm_attn_backend: TRT-LLM 1.3.0rc20 serving forces # attn_backend=FLASHINFER for Gemma4 on every SM # (Gemma4ForConditionalGeneration.get_model_defaults, @@ -42,7 +46,7 @@ model_case_values: window_size: 1024 sglang_runtime_window_size: 1023 sglang_scaling: 1.0 - sglang_backends: {90: triton, 100: trtllm_mha, 103: trtllm_mha, 120: triton} + sglang_backends: {89: triton, 90: triton, 100: trtllm_mha, 103: trtllm_mha, 120: triton} trtllm_attn_backend: FLASHINFER tensor_parallel_sizes: [1, 2, 4, 8, 16] - id: gemma4_global_attention @@ -54,7 +58,7 @@ model_case_values: head_dim: 512 window_size: 0 sglang_scaling: 1.0 - sglang_backends: {90: triton, 100: trtllm_mha, 103: trtllm_mha, 120: triton} + sglang_backends: {89: triton, 90: triton, 100: trtllm_mha, 103: trtllm_mha, 120: triton} trtllm_attn_backend: FLASHINFER tensor_parallel_sizes: [1, 2, 4, 8, 16] - id: gemma4_31b_swa_attention @@ -65,7 +69,7 @@ model_case_values: window_size: 1024 sglang_runtime_window_size: 1023 sglang_scaling: 1.0 - sglang_backends: {90: triton, 100: trtllm_mha, 103: trtllm_mha, 120: triton} + sglang_backends: {89: triton, 90: triton, 100: trtllm_mha, 103: trtllm_mha, 120: triton} trtllm_attn_backend: FLASHINFER tensor_parallel_sizes: [1, 2, 4, 8, 16, 32] - id: gemma4_31b_global_attention @@ -75,7 +79,7 @@ model_case_values: head_dim: 512 window_size: 0 sglang_scaling: 1.0 - sglang_backends: {90: triton, 100: trtllm_mha, 103: trtllm_mha, 120: triton} + sglang_backends: {89: triton, 90: triton, 100: trtllm_mha, 103: trtllm_mha, 120: triton} trtllm_attn_backend: FLASHINFER tensor_parallel_sizes: [1, 2, 4, 8, 16, 32] moe: diff --git a/python/aisimulate/collector/cases/models/Glm5NextForConditionalGeneration_cases.yaml b/python/aisimulate/collector/cases/models/Glm5NextForConditionalGeneration_cases.yaml index 52ae01e1a0..68ed3fa4ba 100644 --- a/python/aisimulate/collector/cases/models/Glm5NextForConditionalGeneration_cases.yaml +++ b/python/aisimulate/collector/cases/models/Glm5NextForConditionalGeneration_cases.yaml @@ -24,8 +24,9 @@ description: > Checkpoints: zai-org/GLM-5.3-Flash (FP8 e4m3 128x128 block for dense MLP, shared + routed experts and MLA q_a/q_b/kv_a/o; BF16 elsewhere) and nvidia/GLM-5.3-Flash-NVFP4 (NVFP4 routed experts + dense MLP; everything - else, including the shared expert, BF16). Runtimes are model-pinned in - framework_manifest.yaml (stock vLLM 0.31.0, SGLang 0.5.20). + else, including the shared expert, BF16). KDA / mHC are audited on stock + vLLM 0.31.0 only: the SGLang 0.5.20 lane retired (#649), and 0.5.21 needs a + port of both paths before SGLang collects them again (the 0.5.20 tables stay). Generic ops land in the generic tables (gemm_perf, moe_perf, kda_perf, mhc_module_perf, comm); only the GLM attention layer gets a GLM-specific module table (DeepSeek-V4.1 precedent). @@ -204,8 +205,6 @@ framework_specific_op_cases: sglang: moe: cases: all - # SGLang 0.5.20 is the audited GLM KDA / mHC runtime. - kda: - cases: all - mhc_module: - cases: all + # No kda / mhc_module: SGLang 0.5.21 is not an audited GLM KDA / mHC runtime + # (collect_kda.py _KDA_ARCHITECTURE_COMPAT, collect_mhc_module.py + # _ARCHITECTURE_COMPAT); scheduling them would only log certain failures. diff --git a/python/aisimulate/collector/cases/models/GlmMoeDsaForCausalLM_cases.yaml b/python/aisimulate/collector/cases/models/GlmMoeDsaForCausalLM_cases.yaml index dca784c116..00ac121374 100644 --- a/python/aisimulate/collector/cases/models/GlmMoeDsaForCausalLM_cases.yaml +++ b/python/aisimulate/collector/cases/models/GlmMoeDsaForCausalLM_cases.yaml @@ -16,6 +16,8 @@ model_paths: - nvidia/GLM-5.2-NVFP4 # GLM-5.3: the plain zai-org repository is the FP8-block release and the # BF16 weights carry the -BF16 suffix (no zai-org/GLM-5.3-FP8 exists). + # Do not spell the non-existent name here: the roster derives the coverage + # floor from repo mentions in these files. - zai-org/GLM-5.3 - zai-org/GLM-5.3-BF16 - nvidia/GLM-5.3-NVFP4 @@ -207,11 +209,15 @@ model_case_values: - zai-org/GLM-5.2 - zai-org/GLM-5.2-FP8 - nvidia/GLM-5.2-NVFP4 - - zai-org/GLM-5.3 - zai-org/GLM-5.3-BF16 + - zai-org/GLM-5.3 - nvidia/GLM-5.3-NVFP4 attention_type: dsa native_num_heads: 64 + # Real tensor-parallel shards (model-correlated narrowing): per-rank + # heads = 64 // tp (64/32/16/8/4); the shared inner head grid's 128 and + # 96/48/24/12 entries are not GLM-5.x shards. + tensor_parallel_sizes: [1, 2, 4, 8, 16] framework_specific_model_case_values: vllm: @@ -226,6 +232,7 @@ framework_specific_model_case_values: architecture: GlmMoeDsaForCausalLM attention_type: dsa native_num_heads: 64 + tensor_parallel_sizes: [1, 2, 4, 8, 16] all_frameworks_op_cases: gemm: @@ -239,8 +246,6 @@ framework_specific_op_cases: cases: all dsa_generation_module: cases: all - moe_ep: - cases: all vllm: dsa_context_module: cases: all diff --git a/python/aisimulate/collector/cases/models/KimiK25ForConditionalGeneration_cases.yaml b/python/aisimulate/collector/cases/models/KimiK25ForConditionalGeneration_cases.yaml index c0a920a7b4..fcae8120c7 100644 --- a/python/aisimulate/collector/cases/models/KimiK25ForConditionalGeneration_cases.yaml +++ b/python/aisimulate/collector/cases/models/KimiK25ForConditionalGeneration_cases.yaml @@ -213,8 +213,6 @@ framework_specific_op_cases: cases: all mla_bmm_gen_post: cases: all - moe_ep: - cases: all sglang: mla_context: cases: all diff --git a/python/aisimulate/collector/cases/models/Llama4ForConditionalGeneration_cases.yaml b/python/aisimulate/collector/cases/models/Llama4ForConditionalGeneration_cases.yaml index 0292d8963b..d0c760befd 100644 --- a/python/aisimulate/collector/cases/models/Llama4ForConditionalGeneration_cases.yaml +++ b/python/aisimulate/collector/cases/models/Llama4ForConditionalGeneration_cases.yaml @@ -39,7 +39,7 @@ model_case_values: window_sizes: [0, 8192] sglang_runtime_window_size: -1 sglang_attention_chunk_size: 8192 - sglang_backends: {90: fa3, 100: trtllm_mha, 103: trtllm_mha, 120: triton} + sglang_backends: {89: triton, 90: fa3, 100: trtllm_mha, 103: trtllm_mha, 120: triton} # sm89: arg_groups/model_overrides/llama4.py "other platforms" -> triton tensor_parallel_sizes: [1, 2, 4, 8] moe: - model_path: meta-llama/Llama-4-Scout-17B-16E-Instruct diff --git a/python/aisimulate/collector/cases/models/MiniMaxM2ForCausalLM_cases.yaml b/python/aisimulate/collector/cases/models/MiniMaxM2ForCausalLM_cases.yaml index 2ca036300a..ff7d15da02 100644 --- a/python/aisimulate/collector/cases/models/MiniMaxM2ForCausalLM_cases.yaml +++ b/python/aisimulate/collector/cases/models/MiniMaxM2ForCausalLM_cases.yaml @@ -74,6 +74,3 @@ all_frameworks_op_cases: cases: all framework_specific_op_cases: - trtllm: - moe_ep: - cases: all diff --git a/python/aisimulate/collector/cases/models/MiniMaxM3ForCausalLM_cases.yaml b/python/aisimulate/collector/cases/models/MiniMaxM3ForCausalLM_cases.yaml index 546076296a..e7e852a0cd 100644 --- a/python/aisimulate/collector/cases/models/MiniMaxM3ForCausalLM_cases.yaml +++ b/python/aisimulate/collector/cases/models/MiniMaxM3ForCausalLM_cases.yaml @@ -87,6 +87,13 @@ model_case_values: - nvidia/MiniMax-M3-NVFP4 attention_type: msa native_num_heads: 64 + # Real shards only (owner decision 2026-09-29): the MSA collectors + # derive per-rank heads as 64 // tp -> 64/32/16/8/4 (kv/index heads + # 4/2/1/1/1). The shared 12-value head grid also produced 128/96/48/ + # 24/12, which no M3 deployment reaches; 48 and 96 shard to 3/6 index + # heads and only ever failed the decode scorer's power-of-2 tile. + # TP 16 is the ceiling anyone runs this model at. + tensor_parallel_sizes: [1, 2, 4, 8, 16] all_frameworks_op_cases: gemm: diff --git a/python/aisimulate/collector/collect.py b/python/aisimulate/collector/collect.py index cdcd958c04..175d2fc515 100644 --- a/python/aisimulate/collector/collect.py +++ b/python/aisimulate/collector/collect.py @@ -11,10 +11,11 @@ import contextlib import functools +import logging import os import warnings -from helper import get_device_module, get_device_str +from helper import cuda_context_error, get_device_module, get_device_str def setup_warning_filters(): @@ -89,6 +90,7 @@ def setup_warning_filters(): from pathlib import Path from helper import ( + EXIT_CODE_OPS_NOT_IN_PLAN, EXIT_CODE_RESTART, PERF_TRANSACTION_FILENAME, SIDECAR_STAGING_FILENAME, @@ -115,14 +117,45 @@ def setup_warning_filters(): perf_transaction_artifact_paths, restore_perf_publications, save_error_report, + set_worker_heartbeat, setup_logging, setup_signal_handlers, validate_restored_perf_publications, ) -logger = None +# Module logger. ``setup_logging`` (main / workers) installs the configured +# collector logger; until then this is the bare library logger so in-process +# callers (fullnode finalization, unit tests, tooling) that reach the checkpoint +# and provenance paths never trip on a None logger. +logger = logging.getLogger("collector") +_LOGGER_CONFIGURED = False RESUME_SCHEMA_VERSION = "collector-resume-v2" STALL_THRESHOLD = 30 # iterations (x 0.5 s sleep = 15 s) before logging a stall warning +# Seconds of ZERO progress with live workers before the stall watchdog acts. +# Progress = a task finished, or a worker heartbeat (helper.worker_heartbeat, +# one per benchmark_with_power call): one task measuring hundreds of points +# stays alive, a point wedged inside a kernel does not. +# dump every worker's thread stacks (SIGUSR1 -> faulthandler), kill the workers, +# record their in-flight tasks as WorkerStallTimeout and restart the workers so +# the queue keeps draining. 0 disables. sglang 0.5.21 moe/bf16 on l40s/b300 +# (2026-10-06) sat 76 min at 0/58482 with no error until the pipeline watchdog +# cancelled the whole job; a finite, diagnosable failure is worth more than an +# open-ended wait. +STALL_TIMEOUT_SEC = int(os.environ.get("AISIM_STALL_TIMEOUT_SEC") or os.environ.get("AIC_STALL_TIMEOUT_SEC") or 1800) +# Consecutive stall rounds (kill + restart) with NO task finished by a worker in +# between before the executor closes the op (abandons the worker slots; the +# never-dispatched tasks are recorded as WorkerOrphanedTask, not as failed). +# 1 = close on the first stall. b200 vllm dsa_context_module s256 (job +# 481670367, 2026-10-11): the first stall at 13,002/21,960 — eight workers each +# wedged on a poison shape — closed the op and wrote the 8,950 undispatched +# tasks into the failed ledger, so the retry-failed resume re-ran all of them +# behind the same poison shapes and stalled again; two full loops were burnt. +STALL_MAX_ROUNDS = max(1, int(os.environ.get("AISIM_STALL_MAX_ROUNDS") or 3)) +# Seconds of zero progress after which the executor asks every live worker for a +# stack dump (SIGUSR1 -> faulthandler) ONCE, without killing: the l40s sglang +# moe/bf16 first case legitimately took minutes (JIT) and then the shard ran 2h23m +# to completion — the diagnosis must not cost the run. +STALL_DIAG_SEC = int(os.environ.get("AISIM_STALL_DIAG_SEC") or 600) # Failures of one (model, dtype) group within an op before the summary flags # it as systemic (a fix-me warning; nothing is skipped). SYSTEMIC_GROUP_THRESHOLD = 5 @@ -165,7 +198,7 @@ def setup_warning_filters(): @dataclass(frozen=True) class _ProducerCheckpointPlan: attestation: "_FileAttestation" - table: str + tables: tuple[str, ...] # every registry-declared table this producer owns (sorted) identity: tuple[tuple[str, object], ...] done: frozenset[str] attempted: frozenset[str] @@ -194,7 +227,7 @@ def identity(self) -> tuple[int, int]: @dataclass(frozen=True) class _ValidatedCheckpoint: attestation: _FileAttestation - table: str + tables: tuple[str, ...] attempted: frozenset[str] document: dict @@ -1111,6 +1144,131 @@ def _get_test_cases_for_model(get_func, model_path: str | None): return get_func() +class ReplaySpec: + """Failure replay (``--cases-from``): run exactly the cases an earlier run's failure + records name, plus an optional random control sample per op. + + Selection is by exact case string — the same ``str(task)`` the failure records store + as ``task_params`` — never by substring and never by reconstructing tuples, so the + requested set is unambiguous. Nothing else about the run changes: same getter, same + run_func, same dispatch, same measurement, so a replay observes exactly what the + campaign observes. The run leaves ``collector_replay.json`` in the output root. It is + a diagnosis, not a dataset: main() refuses an output root that already holds results + and never finalizes a replay's CSV staging into parquet or ``collection_meta.yaml``, + so its rows cannot replace or extend a table. This is a selection mechanism like + ``--case-filter``: it decides WHICH planned cases run, never whether a case may be + skipped once queued. + """ + + MARKER = "collector_replay.json" + + def __init__(self, source: str, wanted: set[str], control: int = 0): + self.source = source + self.wanted = set(wanted) + self.control = int(control) + self.matched: dict[str, set[str]] = {} + self.controls: dict[str, int] = {} + + @staticmethod + def case_key(case) -> str: + """The string a failure record stores for this case (see collect_module_safe).""" + if isinstance(case, dict) and "id" in case and "params" in case: + return str(case["params"]) + return str(case) + + def select(self, op: str, cases: list, rng: random.Random) -> list: + wanted = [case for case in cases if self.case_key(case) in self.wanted] + self.matched[op] = {self.case_key(case) for case in wanted} + controls: list = [] + if self.control > 0: + rest = [case for case in cases if self.case_key(case) not in self.wanted] + controls = rng.sample(rest, min(self.control, len(rest))) + self.controls[op] = len(controls) + return wanted + controls + + def unmatched(self) -> list[str]: + seen: set[str] = set().union(*self.matched.values()) if self.matched else set() + return sorted(self.wanted - seen) + + def report(self) -> dict: + return { + "source": self.source, + "requested": len(self.wanted), + "matched": sum(len(v) for v in self.matched.values()), + "matched_by_op": {op: len(v) for op, v in self.matched.items()}, + "controls_by_op": dict(self.controls), + "unmatched": self.unmatched(), + } + + +def load_replay_cases(path: "str | Path") -> set[str]: + """Case strings for ``--cases-from``: an ``errors_*.json[.gz]`` or + ``collection_summary_*.json[.gz]`` written by an earlier run (``task_params`` of every + entry that has one), a JSON list of case strings, or a text file with one case per line.""" + import gzip + + p = Path(path) + if p.suffix == ".gz": + with gzip.open(p, "rt", encoding="utf-8") as fh: + raw = fh.read() + else: + raw = p.read_text(encoding="utf-8") + try: + doc = json.loads(raw) + except json.JSONDecodeError: + doc = [line.strip() for line in raw.splitlines() if line.strip()] + if isinstance(doc, dict): + doc = doc.get("errors", []) + out: set[str] = set() + for item in doc: + if isinstance(item, str): + out.add(item) + elif isinstance(item, dict) and item.get("task_params") not in (None, "None", ""): + out.add(str(item["task_params"])) + if not out: + raise ValueError(f"--cases-from {path}: no case strings found") + return out + + +def _load_replay_spec(args, log) -> "ReplaySpec | None": + if not getattr(args, "cases_from", None): + return None + wanted = load_replay_cases(args.cases_from) + spec = ReplaySpec(str(args.cases_from), wanted, control=args.replay_control or 0) + log.info(f"Replay mode: {len(wanted)} requested cases from {args.cases_from}; control sample {spec.control} per op") + return spec + + +def _require_replay_output_root_without_results( + spec: "ReplaySpec | None", output_root: Path, existing_perf_outputs +) -> None: + """A replay re-measures cases an earlier run already wrote; in that run's output root its + rows would replace or extend the dataset, so it needs a directory without results.""" + if spec is None: + return + if existing_perf_outputs or (output_root / "collection_meta.yaml").exists() or any(output_root.glob("*.parquet")): + raise SystemExit( + f"--cases-from: {output_root} already holds collector results; run a replay in an empty " + "directory so a diagnosis can never overwrite or extend a dataset" + ) + + +def _finish_replay(spec: "ReplaySpec | None", output_root: Path, log) -> None: + if spec is None: + return + report = spec.report() + (output_root / ReplaySpec.MARKER).write_text(json.dumps(report, indent=1) + "\n", encoding="utf-8") + if report["unmatched"]: + log.warning( + f"Replay: {len(report['unmatched'])}/{report['requested']} requested cases are not in this run's " + f"case plan (first: {report['unmatched'][:3]})" + ) + log.info( + f"Replay: matched {report['matched']}/{report['requested']} requested cases; " + f"{ReplaySpec.MARKER} written — this result is a diagnosis, not a dataset" + ) + + def _requested_ops(ops: list[str] | None, case_plan=None) -> set[str]: if ops is not None: return set(ops) @@ -1229,13 +1387,15 @@ def __init__( def load_existing(self): """Load an existing checkpoint for resume. Raises on mismatch.""" if not self._path.parent.exists() and not self._path.parent.is_symlink(): - logger.info(f"{self.module_name}: no checkpoint found, starting fresh") + if logger is not None: + logger.info(f"{self.module_name}: no checkpoint found, starting fresh") return with _locked_checkpoint_directory(self._path) as locked: canonical_path = _checkpoint_path_in_locked_directory(self._path, locked) _normalize_atomic_replace_state_at(locked.file_descriptor, canonical_path) if _entry_state_at(locked.file_descriptor, canonical_path.name) is None: - logger.info(f"{self.module_name}: no checkpoint found, starting fresh") + if logger is not None: + logger.info(f"{self.module_name}: no checkpoint found, starting fresh") return try: @@ -1276,7 +1436,8 @@ def load_existing(self): self._source_digest = snapshot.digest self._source_device = snapshot.device self._source_inode = snapshot.inode - logger.info(f"{self.module_name}: loaded checkpoint — {len(self._done)} passed, {len(self._failed)} failed") + if logger is not None: + logger.info(f"{self.module_name}: loaded checkpoint — {len(self._done)} passed, {len(self._failed)} failed") # -- public API ------------------------------------------------------- @@ -1440,6 +1601,21 @@ def __exit__(self, exc_type, exc_val, exc_tb): logger.info(f"Full profile saved to: {profile_file}") +KERNEL_LIMIT_PREFIX = "FIXME(kernel-limit)" + + +def _classify_exception(error: BaseException) -> str: + """How the error report labels a failed case. + + ``kernel_limit``: the collector itself refused the cell with a cited, + deterministic framework limit (message starts with FIXME(kernel-limit) — + divisibility guards, missing kernels, unsupported head ranges). These are + expected on every rerun and must be counted apart from ``unexpected`` + failures (pipeline summaries matched message text by hand until 2026-10-08). + """ + return "kernel_limit" if str(error).lstrip().startswith(KERNEL_LIMIT_PREFIX) else "unexpected" + + def _failure_group(task) -> str | None: """Group label for failure aggregation: one (model, dtype) family within an op. @@ -1472,6 +1648,11 @@ def _is_cuda_fatal_exception(exc, torch_mod) -> bool: error_text = str(exc).lower() fatal_markers = ( "illegal memory access", + "cudaerrorillegaladdress", + # FlashInfer's TVM-FFI launchers report a sticky error as + # "Error in function 'run' at : Got CUDA error. See above for + # details." -- the error text itself only goes to stderr. + "got cuda error", "unspecified launch failure", "cuda_error_launch_failed", "cublas_status_execution_failed", @@ -1544,6 +1725,7 @@ def worker( module_name="unknown", current_task_ids=None, consumed_sentinel=None, + heartbeat=None, ): """worker with automatic logging setup""" @@ -1565,6 +1747,15 @@ def worker( get_device_module().set_device(device) worker_logger.info(f"Worker {device_id} initialized for {module_name}") + # Liveness for the executor's stall watchdog: every measurement beats through + # helper.worker_heartbeat (benchmark_with_power), so one long task counts as + # progress before it finishes. + def _beat(): + if heartbeat is not None: + heartbeat.value = time.time() + + set_worker_heartbeat(_beat if heartbeat is not None else None) + # Process tasks while True: task_info = queue.get() @@ -1594,6 +1785,7 @@ def worker( try: worker_logger.debug(f"Starting task {task_id}") + _beat() result = func(*task, device=device) # Only the dedicated sentinel requests a recycle: entrypoints also # return plain ints (row counts), which must never be mistaken for @@ -1632,7 +1824,7 @@ def worker( "task_params": str(task), "error_type": type(e).__name__, "error_message": str(e), - "classification": "unexpected", + "classification": _classify_exception(e), "group": _failure_group(task), "traceback": traceback.format_exc(), "timestamp": datetime.now().isoformat(), @@ -1658,7 +1850,18 @@ def worker( for handler in worker_logger.handlers: handler.flush() - if _is_cuda_fatal_exception(e, torch_mod): + # Message matching alone misses launcher wrappers that hide the CUDA + # error text (FlashInfer's "Got CUDA error. See above for details."): + # the worker then runs the next case on a dead context and books an + # unrelated failure for it (gb200 job 512476645, 2026-10-27: one + # healthy case lost per Llama-4 autotune crash). Probe the context + # after every failure and recycle the worker when it is gone. + context_error = None if _is_cuda_fatal_exception(e, torch_mod) else cuda_context_error(device) + if context_error is not None: + worker_logger.warning( + f"CUDA context unusable after task {task_id} failed with {type(e).__name__}: {context_error}" + ) + if context_error is not None or _is_cuda_fatal_exception(e, torch_mod): worker_logger.warning( f"Fatal {type(e).__name__} encountered on task {task_id}. " f"Worker {device_id} exiting to reset GPU context. " @@ -1755,6 +1958,9 @@ def _unresolved_failure_errors(): manager = mp.Manager() progress_value = manager.Value("i", 0) lock = manager.Lock() + # Last time any worker reported liveness (helper.worker_heartbeat); the + # stall watchdog treats a fresh beat like a finished task. + heartbeat = mp.Value("d", 0.0, lock=False) # Track process health process_stats = {i: {"restarts": 0, "errors": []} for i in range(num_processes)} @@ -1788,6 +1994,7 @@ def start_process(device_id): module_name, current_task_ids, consumed_sentinel, + heartbeat, ), ) p.start() @@ -1873,9 +2080,81 @@ def sync_done_to_checkpoint(): # Monitor progress with error collection errors = [] + # Worker slots the stall watchdog killed on purpose. The health check below + # restarts them through start_process without recording a WorkerSignalCrash: + # their in-flight task is already a WorkerStallTimeout failure. + watchdog_killed: set[int] = set() + + def _dump_and_kill_stalled_workers(stalled_for: float, *, stall_round: int, close_op: bool): + """Stall watchdog: stacks first (SIGUSR1 -> faulthandler in the worker), + then SIGTERM/SIGKILL. The in-flight task of every live worker becomes a + WorkerStallTimeout failure; the killed workers are restarted by the + health check and the queue keeps draining. Only when ``close_op`` is set + (STALL_MAX_ROUNDS consecutive rounds with nothing finished in between) + are the slots abandoned so the orphan path below closes the op.""" + live = [(i, p) for i, p in enumerate(processes) if p is not None and p.is_alive()] + action = "closing the op" if close_op else "restarting the workers" + logger.error( + f"{module_name}: no progress for {int(stalled_for)} s with {len(live)} live worker(s) " + f"(STALL_TIMEOUT_SEC={STALL_TIMEOUT_SEC}, stall round {stall_round}/{STALL_MAX_ROUNDS}); " + f"dumping worker stacks, failing the in-flight tasks and {action}" + ) + for _, p in live: + try: + os.kill(p.pid, signal.SIGUSR1) + except Exception: + pass + time.sleep(5) # let faulthandler write before the kill + for i, p in live: + active_task_id = current_task_ids.get(i) + if active_task_id is not None and active_task_id not in accounted: + # same task_params string a worker-side failure stores, so --cases-from can replay it + params = next((t["params"] for t in task_infos if t["id"] == active_task_id), None) + errors.append( + { + "module": module_name, + "device_id": i, + "task_id": active_task_id, + "task_params": None if params is None else str(params), + "error_type": "WorkerStallTimeout", + "error_message": f"worker made no progress for {int(stalled_for)} s; stacks dumped via SIGUSR1", + "classification": "unexpected", + "group": None if params is None else _failure_group(params), + "traceback": "", + "timestamp": datetime.now().isoformat(), + } + ) + try: + failed_tasks[active_task_id] = True + except Exception: + pass + try: + current_task_ids[i] = None + except Exception: + pass + p.terminate() + p.join(timeout=10) + if p.is_alive(): + p.kill() + p.join(timeout=10) + if p.is_alive(): + # wedged beyond SIGKILL (driver call): its GPU cannot serve a + # replacement worker either, so the slot is abandoned + logger.error(f"Process {i} ({p.pid}) survived SIGKILL; abandoning its worker slot") + processes[i] = None + elif close_op: + processes[i] = None + else: + watchdog_killed.add(i) + with tqdm(total=len(task_infos), desc=f"{module_name}", dynamic_ncols=True, leave=True) as pbar: last_progress = 0 + last_heartbeat = heartbeat.value stall_count = 0 + stall_since = None + stall_diag_sent = False + stall_rounds = 0 + accounted_after_stall_kill = None last_error_count = 0 if num_processes == 0: @@ -1904,7 +2183,7 @@ def sync_done_to_checkpoint(): "task_params": str(task_params), "error_type": type(e).__name__, "error_message": str(e), - "classification": "unexpected", + "classification": _classify_exception(e), "group": _failure_group(task_params), "traceback": traceback.format_exc(), "timestamp": datetime.now().isoformat(), @@ -1932,14 +2211,65 @@ def sync_done_to_checkpoint(): pbar.set_postfix({"errors": len(errors)}) last_error_count = len(errors) - if len(accounted) == last_progress: + if len(accounted) == last_progress and heartbeat.value > last_heartbeat: + # No task finished, but a worker measured another point of a long + # task since the last check (helper.worker_heartbeat): progress. + last_heartbeat = heartbeat.value + stall_count = 0 + stall_since = None + stall_diag_sent = False + elif len(accounted) == last_progress: stall_count += 1 + stall_since = stall_since or time.time() if stall_count > STALL_THRESHOLD: - logger.warning(f"Progress stalled at {len(accounted)}/{len(task_infos)}") + if not accounted: + # the first case carries module import + JIT/autotune: slow is normal, silent is not + logger.info( + f"waiting for the first result ({int(time.time() - stall_since)} s, 0/{len(task_infos)}; " + f"stacks dumped at {STALL_DIAG_SEC} s, workers restarted at {STALL_TIMEOUT_SEC} s of no " + f"progress, op stopped after {STALL_MAX_ROUNDS} such rounds without progress)" + ) + else: + logger.warning(f"Progress stalled at {len(accounted)}/{len(task_infos)}") stall_count = 0 + if STALL_DIAG_SEC > 0 and not stall_diag_sent and time.time() - stall_since > STALL_DIAG_SEC: + stall_diag_sent = True + live = [p for p in processes if p is not None and p.is_alive()] + logger.warning( + f"{module_name}: no progress for {int(time.time() - stall_since)} s; requesting thread " + f"stacks from {len(live)} live worker(s) (SIGUSR1 -> faulthandler, workers keep running)" + ) + for p in live: + try: + os.kill(p.pid, signal.SIGUSR1) + except Exception: + pass + if ( + STALL_TIMEOUT_SEC > 0 + and time.time() - stall_since > STALL_TIMEOUT_SEC + and any(p is not None and p.is_alive() for p in processes) + ): + # Rounds are consecutive only while no task has been finished + # (by a worker) since the previous round's kill. + if accounted_after_stall_kill is not None and len(accounted) == accounted_after_stall_kill: + stall_rounds += 1 + else: + stall_rounds = 1 + _dump_and_kill_stalled_workers( + time.time() - stall_since, + stall_round=stall_rounds, + close_op=stall_rounds >= STALL_MAX_ROUNDS, + ) + sync_done_to_checkpoint() # the WorkerStallTimeout tasks are accounted from here on + accounted_after_stall_kill = len(accounted) + stall_since = None + stall_diag_sent = False else: stall_count = 0 + stall_since = None + stall_diag_sent = False last_progress = len(accounted) + last_heartbeat = heartbeat.value # Check process health — only restart if there is still work # remaining. Workers that consumed a None sentinel or finished @@ -1985,7 +2315,13 @@ def sync_done_to_checkpoint(): with lock: progress_value.value += 1 - crash_error = create_process_exit_error(i, exit_code) + if i in watchdog_killed: + # killed by the stall watchdog on purpose: the in-flight + # task is already a WorkerStallTimeout failure + watchdog_killed.discard(i) + crash_error = None + else: + crash_error = create_process_exit_error(i, exit_code) if crash_error: errors.append(crash_error) process_stats[i]["errors"].append("process_exit") @@ -2008,13 +2344,18 @@ def sync_done_to_checkpoint(): processes[i] = None # Escape hatch: every worker slot is permanently gone (restart - # limit, consumed sentinel, or no remaining work to justify a - # restart) while tasks are still unaccounted — e.g. a worker was - # SIGKILLed after queue.get() but before it recorded anything in - # done_tasks/current_task_ids. Without this the monitor loop - # would spin at 0.5s forever with nothing able to make progress. - # Record the orphaned ids as failures so the summary and resume - # checkpoint stay complete, then stop the loop. + # limit, consumed sentinel, stall watchdog closing the op, or no + # remaining work to justify a restart) while tasks are still + # unaccounted — e.g. a worker was SIGKILLed after queue.get() but + # before it recorded anything in done_tasks/current_task_ids. + # Without this the monitor loop would spin at 0.5s forever with + # nothing able to make progress. The orphaned ids are recorded in + # the error report (WorkerOrphanedTask) for the summary and the + # harness, but NOT in the failed ledger: nothing ran them, so a + # plain --resume must queue them again. 2026-10-11, b200 vllm + # dsa_context_module: 8,950 orphans written as failed were re-run + # by --resume-retry-failed behind the poison shapes that had + # stalled the op, which stalled it again. if all(p is None for p in processes) and len(accounted) < len(task_infos): while not error_queue.empty(): error = error_queue.get() @@ -2025,7 +2366,8 @@ def sync_done_to_checkpoint(): if orphaned: logger.error( f"All workers exited with {len(orphaned)} task(s) unaccounted; " - f"recording them as failed (orphaned by worker death) and stopping the monitor loop." + f"recording them as WorkerOrphanedTask (never dispatched — left out of the failed " + f"ledger so a plain --resume runs them) and stopping the monitor loop." ) for task_id in orphaned: errors.append( @@ -2035,14 +2377,14 @@ def sync_done_to_checkpoint(): "task_id": task_id, "task_params": None, "error_type": "WorkerOrphanedTask", - "error_message": "all workers exited before this task was accounted", + "error_message": "all workers exited before this task was dispatched; " + "not recorded as failed — a plain --resume queues it again", "classification": "unexpected", "group": None, "traceback": "", "timestamp": datetime.now().isoformat(), } ) - failed_tasks[task_id] = True sync_done_to_checkpoint() break @@ -2068,6 +2410,14 @@ def sync_done_to_checkpoint(): if p.is_alive(): logger.warning(f"Process {p.pid} did not terminate, forcing...") p.terminate() + p.join(timeout=10) + if p.is_alive(): + # a worker wedged in a CUDA call ignores SIGTERM; without SIGKILL the + # interpreter's multiprocessing atexit join waits forever (pipeline + # watchdogs had to reap collect.py itself, 2026-10-05 campaign) + logger.warning(f"Process {p.pid} survived SIGTERM, killing") + p.kill() + p.join(timeout=10) # Shutdown manager to clean up resources (semaphores, etc.) manager.shutdown() @@ -2107,6 +2457,7 @@ def collect_ops( case_plan=None, sm_version: int | None = None, case_filters: list[str] | None = None, + replay: ReplaySpec | None = None, ) -> list[dict]: """Run collection for a list of resolved collection entries. @@ -2176,11 +2527,15 @@ class CompatibilityError(RuntimeError): get_func = getattr(get_module, collection["get_func"]) run_func = getattr(run_module, collection["run_func"]) - perf_filename = collection["perf_filename"] - worker_filename = collection.get("worker_perf_filename") - if worker_filename is not None: - perf_filename = str(Path(perf_filename).with_name(worker_filename)) - run_func = functools.partial(run_func, perf_filename=perf_filename) + run_func = functools.partial(run_func, perf_filename=collection["perf_filename"]) + if collection.get("extra_perf_filenames"): + # a multi-table producer (registry OpEntry.extra_perf_filenames) receives + # every table it must write; finalize binds all of them to its checkpoint + run_func = functools.partial(run_func, extra_perf_filenames=tuple(collection["extra_perf_filenames"])) + if collection.get("run_kwargs"): + # registry OpEntry.run_kwargs: fixed mode arguments for a run_func shared by + # several ops (sglang dsa_*_module_skip_indexer -> skip_indexer=True) + run_func = functools.partial(run_func, **dict(collection["run_kwargs"])) def get_func_with_limit(get_func=get_func, op=collection["type"]): from collector.capabilities import filter_cases @@ -2191,6 +2546,13 @@ def get_func_with_limit(get_func=get_func, op=collection["type"]): before_count = len(cases) cases = [case for case in cases if any(fragment in str(case) for fragment in case_filters)] logger.info(f"{op}: --case-filter kept {len(cases)}/{before_count} cases") + if replay is not None: + before_count = len(cases) + cases = replay.select(op, cases, random.Random(shuffle_seed)) + logger.info( + f"{op}: replay selected {len(replay.matched[op])} requested + " + f"{replay.controls[op]} control of {before_count} cases" + ) if shuffle: rng = random.Random(shuffle_seed) rng.shuffle(cases) @@ -2238,6 +2600,7 @@ def collect_sglang( case_plan=None, sm_version: int | None = None, case_filters: list[str] | None = None, + replay: ReplaySpec | None = None, ): """Collect performance data for SGLang with enhanced error tracking""" os.environ["FLASHINFER_LOG_LEVEL"] = "ERROR" @@ -2297,6 +2660,7 @@ def collect_sglang( case_plan=case_plan, sm_version=sm_version, case_filters=case_filters, + replay=replay, ) for collection in fullnode_collections: @@ -2313,6 +2677,7 @@ def collect_sglang( case_plan=case_plan, sm_version=sm_version, case_filters=case_filters, + replay=replay, get_test_cases_for_model=_get_test_cases_for_model, resume_checkpoint_cls=ResumeCheckpoint, logger=logger, @@ -2340,6 +2705,7 @@ def collect_vllm( case_plan=None, sm_version: int | None = None, case_filters: list[str] | None = None, + replay: ReplaySpec | None = None, ): """Collect performance data for vLLM""" from collector.version_resolver import build_collections @@ -2397,6 +2763,7 @@ def collect_vllm( case_plan=case_plan, sm_version=sm_version, case_filters=case_filters, + replay=replay, ) generate_collection_summary(all_errors, "vllm", version) @@ -2420,6 +2787,7 @@ def collect_trtllm( case_plan=None, sm_version: int | None = None, case_filters: list[str] | None = None, + replay: ReplaySpec | None = None, ): """Collect performance data for TensorRT LLM with enhanced error tracking""" from collector.trtllm.registry import REGISTRY @@ -2468,6 +2836,7 @@ def collect_trtllm( case_plan=case_plan, sm_version=sm_version, case_filters=case_filters, + replay=replay, ) generate_collection_summary(all_errors, "trtllm", version) @@ -2569,7 +2938,15 @@ def _all_op_names() -> list[str]: def _git_collector_ref(repo_root: Path) -> str: - """The repo SHA the collector ran from (design §5), "unknown" outside a repo.""" + """The repo SHA the collector ran from (design §5), "unknown" outside a repo. + + ``AISIM_COLLECTOR_REF`` (legacy ``AIC_COLLECTOR_REF``) wins when set: the + framework images carry no git, so pipelines that check the collector out + by SHA pass it in (every b200 shard run of 2026-10-05 recorded 'unknown'). + """ + forced = (os.environ.get("AISIM_COLLECTOR_REF") or os.environ.get("AIC_COLLECTOR_REF") or "").strip() + if forced: + return forced try: result = subprocess.run( ["git", "rev-parse", "HEAD"], @@ -2731,8 +3108,24 @@ def _resume_tracker_for_collection( ) -def _registered_checkpoint_table(identity: dict, *, backend: str) -> str: - """Resolve a checkpoint producer to its registry-owned table.""" +def _collection_perf_filenames(collection: dict) -> list[str]: + """Every staged table a provenance collection entry declares (primary first).""" + return [str(collection["perf_filename"]), *(str(name) for name in collection.get("extra_perf_filenames") or ())] + + +def _collection_tables(collection: dict) -> tuple[str, ...]: + return tuple(sorted(Path(name).stem for name in _collection_perf_filenames(collection))) + + +def _registered_checkpoint_tables(identity: dict, *, backend: str) -> frozenset[str]: + """Resolve a checkpoint producer to the set of registry-owned tables it writes. + + One producer (module + run_func) owns one or more tables — the primary + ``perf_filename`` plus any ``extra_perf_filenames`` its OpEntry declares + (owner decision 2026-09-28: compute_scale writes computescale_perf and + scale_matrix_perf from one measurement, so the executor binds both to the + same checkpoint instead of splitting the op). + """ from collector.version_resolver import resolve_module if set(identity) != set(_CHECKPOINT_IDENTITY_FIELDS) or identity.get("backend") != backend: @@ -2756,16 +3149,16 @@ def _registered_checkpoint_table(identity: dict, *, backend: str) -> str: registry = list(registry_module.REGISTRY_XPU) else: registry = [*registry_module.REGISTRY, *_wideep_registry_for_backend(backend)] - owned_tables = { - Path(str(entry.perf_filename)).stem + owned_table_sets = { + frozenset(Path(str(name)).stem for name in entry.perf_filenames) for entry in registry if identity["module"] == f"{backend}.{entry.op}" and identity["run_func"] == entry.run_func and resolve_module(entry, framework_version) is not None } - if len(owned_tables) != 1: - raise RuntimeError(f"checkpoint producer has no unambiguous registered table: {identity!r}") - return owned_tables.pop() + if len(owned_table_sets) != 1: + raise RuntimeError(f"checkpoint producer has no unambiguous registered table set: {identity!r}") + return owned_table_sets.pop() def _load_selected_producer_checkpoint( @@ -2796,14 +3189,14 @@ def _load_selected_producer_checkpoint( f"{context}: staged table {staging_path} has no checkpoint for selected producer " f"{resume_tracker.module_name} at {checkpoint_path}" ) - if _registered_checkpoint_table(resume_tracker._metadata, backend=backend) != staging_path.stem: + if staging_path.stem not in _registered_checkpoint_tables(resume_tracker._metadata, backend=backend): raise RuntimeError( f"{context}: checkpoint producer {resume_tracker.module_name} does not own staged table {staging_path}" ) return resume_tracker -def _producer_checkpoint_plan(resume_tracker: ResumeCheckpoint, table: str) -> _ProducerCheckpointPlan: +def _producer_checkpoint_plan(resume_tracker: ResumeCheckpoint, tables: Iterable[str]) -> _ProducerCheckpointPlan: if ( resume_tracker._source_digest is None or resume_tracker._source_device is None @@ -2817,7 +3210,7 @@ def _producer_checkpoint_plan(resume_tracker: ResumeCheckpoint, table: str) -> _ device=resume_tracker._source_device, inode=resume_tracker._source_inode, ), - table=table, + tables=tuple(sorted(tables)), identity=tuple((field, resume_tracker._metadata[field]) for field in _CHECKPOINT_IDENTITY_FIELDS), done=frozenset(resume_tracker._done), attempted=frozenset(resume_tracker._attempted), @@ -2845,10 +3238,9 @@ def _revalidate_producer_plan(producer_plan: dict[Path, _ProducerCheckpointPlan] ) checkpoint = json.loads(snapshot.contents) identity = {field: checkpoint.get(field) for field in _CHECKPOINT_IDENTITY_FIELDS} - if ( - identity != plan.identity_dict() - or _registered_checkpoint_table(identity, backend=identity["backend"]) != plan.table - ): + if identity != plan.identity_dict() or _registered_checkpoint_tables( + identity, backend=identity["backend"] + ) != frozenset(plan.tables): raise RuntimeError("checkpoint producer identity changed after preflight") if ( _checkpoint_case_ids(checkpoint, "done") != set(plan.done) @@ -2881,16 +3273,18 @@ def _pending_resume_perf_outputs( """ producers_by_output: dict[Path, list[dict]] = {} for collection in provenance_ctx.get("collections") or []: - perf_path = Path(str(collection["perf_filename"])) - if not perf_path.is_absolute(): - perf_path = output_root / perf_path - if perf_path.name.endswith("_perf.txt"): - producers_by_output.setdefault(perf_path, []).append(collection) + for perf_filename in _collection_perf_filenames(collection): + perf_path = Path(perf_filename) + if not perf_path.is_absolute(): + perf_path = output_root / perf_path + if perf_path.name.endswith("_perf.txt"): + producers_by_output.setdefault(perf_path, []).append(collection) pending_outputs: set[Path] = set() for perf_path, producers in producers_by_output.items(): staging_present = perf_path.exists() or perf_path.is_symlink() has_pending_attempts = False + has_completed_cases = False for collection in producers: resume_tracker = _load_selected_producer_checkpoint( collection, @@ -2903,8 +3297,24 @@ def _pending_resume_perf_outputs( context="resume finalization", ) has_pending_attempts = has_pending_attempts or bool(resume_tracker and resume_tracker._attempted) + # _done is cumulative across invocations; only a case that passed inside the open + # event (attempted there) can have staged rows + has_completed_cases = has_completed_cases or bool( + resume_tracker and resume_tracker._done & resume_tracker._attempted + ) if has_pending_attempts: if perf_path.is_symlink() or not perf_path.is_file(): + if not has_completed_cases: + # Every attempted case of every producer failed (e.g. a lane guard + # rejecting the whole shard: sglang moe/int4_wo on 0.5.21/SM100, + # job 469988017): there is legitimately nothing staged. The failures + # are already in the error report; do not turn them into a + # finalization crash that also hides the other tables of the run. + logger.warning( + f"resume finalization: {perf_path} has an open checkpoint event but no completed case " + "and no staging table — all attempted cases failed; nothing to finalize for it" + ) + continue raise RuntimeError( f"resume finalization: open checkpoint event has no regular staging table {perf_path}" ) @@ -3042,10 +3452,11 @@ def _preflight_collector_finalization_inputs( } producers_by_path: dict[Path, list[dict]] = {} for collection in provenance_ctx.get("collections") or []: - staging_path = Path(str(collection["perf_filename"])) - if not staging_path.is_absolute(): - staging_path = output_root / staging_path - producers_by_path.setdefault(staging_path, []).append(collection) + for perf_filename in _collection_perf_filenames(collection): + staging_path = Path(perf_filename) + if not staging_path.is_absolute(): + staging_path = output_root / staging_path + producers_by_path.setdefault(staging_path, []).append(collection) producer_plan: dict[Path, _ProducerCheckpointPlan] = {} seen_attempted_case_ids: set[str] = set() @@ -3067,7 +3478,7 @@ def _preflight_collector_finalization_inputs( ) if resume_tracker is not None: checkpoint_path = resume_tracker._path - checkpoint_plan = _producer_checkpoint_plan(resume_tracker, staging_path.stem) + checkpoint_plan = _producer_checkpoint_plan(resume_tracker, _collection_tables(collection)) prior_plan = producer_plan.get(checkpoint_path) if prior_plan is None: duplicate_case_ids = set(checkpoint_plan.attempted) & seen_attempted_case_ids @@ -3234,7 +3645,7 @@ def _perf_checkpoint_record(record: _ProducerCheckpointPlan) -> dict: "digest": record.attestation.digest, "device": record.attestation.device, "inode": record.attestation.inode, - "table": record.table, + "tables": list(record.tables), "done": sorted(record.done), "failed": sorted(record.failed), "attempted": sorted(record.attempted), @@ -3292,7 +3703,7 @@ def _validate_perf_transaction_document( "digest", "device", "inode", - "table", + "tables", "done", "failed", "attempted", @@ -3300,6 +3711,14 @@ def _validate_perf_transaction_document( } if not isinstance(checkpoint, dict) or set(checkpoint) != expected_fields: raise RuntimeError(f"Invalid collector perf checkpoint record in {journal_path}") + recorded_tables = checkpoint["tables"] + if ( + not isinstance(recorded_tables, list) + or not recorded_tables + or any(not isinstance(table, str) or not table for table in recorded_tables) + or len(recorded_tables) != len(set(recorded_tables)) + ): + raise RuntimeError(f"Invalid collector perf checkpoint tables in {journal_path}") attestation = _attestation_from_record( {field: checkpoint[field] for field in ("path", "digest", "device", "inode")}, None, @@ -3310,8 +3729,8 @@ def _validate_perf_transaction_document( if not isinstance(identity, dict) or set(identity) != set(_CHECKPOINT_IDENTITY_FIELDS): raise RuntimeError(f"Invalid collector perf checkpoint identity in {journal_path}") expected_path = _checkpoint_path(checkpoint_root, identity["module"]) - table = _registered_checkpoint_table(identity, backend=backend) - if attestation.path != expected_path or checkpoint["table"] != table or table not in tables: + owned_tables = _registered_checkpoint_tables(identity, backend=backend) + if attestation.path != expected_path or set(recorded_tables) != owned_tables or not owned_tables <= tables: raise RuntimeError(f"Unowned collector perf checkpoint in {journal_path}: {attestation.path}") if attestation.path in seen_checkpoint_paths: raise RuntimeError(f"Duplicate collector perf checkpoint in {journal_path}: {attestation.path}") @@ -3329,10 +3748,10 @@ def _validate_perf_transaction_document( if attempted & seen_attempted or (not attempted and not (checkpoint["done"] or checkpoint["failed"])): raise RuntimeError(f"Invalid collector perf checkpoint attempts in {journal_path}") if attempted: - attempted_tables.add(table) + attempted_tables.update(owned_tables) seen_checkpoint_paths.add(attestation.path) seen_attempted.update(attempted) - if {checkpoint["table"] for checkpoint in checkpoint_records} != tables: + if {table for checkpoint in checkpoint_records for table in checkpoint["tables"]} != tables: raise RuntimeError(f"Collector perf transaction tables lack checkpoint owners in {journal_path}") if attempted_tables != tables: raise RuntimeError(f"Collector perf transaction tables lack attempted checkpoint case IDs in {journal_path}") @@ -3967,7 +4386,7 @@ def _validated_transaction_checkpoints( raise ValueError("attempted case IDs must be unique across checkpoint participants") if not isinstance(identity, dict) or set(identity) != set(_CHECKPOINT_IDENTITY_FIELDS): raise TypeError(f"checkpoint identity must contain exactly {_CHECKPOINT_IDENTITY_FIELDS!r}") - table = _registered_checkpoint_table(identity, backend=backend) + tables = tuple(sorted(_registered_checkpoint_tables(identity, backend=backend))) checkpoint_path = Path(path_text) expected_path = _checkpoint_path(checkpoint_root, identity["module"]) if path_text != str(expected_path) or checkpoint_path in seen_checkpoint_paths: @@ -3976,12 +4395,12 @@ def _validated_transaction_checkpoints( raise RuntimeError(f"Invalid checkpoint participant in {journal_path}: {error}") from error seen_checkpoint_paths.add(checkpoint_path) seen_attempted_case_ids.update(attempted_case_ids) - parsed_participants.append((checkpoint_path, identity, recorded_ledgers, attempted_case_ids, table)) + parsed_participants.append((checkpoint_path, identity, recorded_ledgers, attempted_case_ids, tables)) validated_participants: list[_ValidatedCheckpoint] = [] with _locked_checkpoint_root(checkpoint_root) as locked: preflight_participants = [] - for checkpoint_path, identity, recorded_ledgers, attempted_case_ids, table in parsed_participants: + for checkpoint_path, identity, recorded_ledgers, attempted_case_ids, tables in parsed_participants: try: canonical_path = _checkpoint_path_in_locked_directory(checkpoint_path, locked) effective_state = _effective_checkpoint_state_at( @@ -4008,14 +4427,14 @@ def _validated_transaction_checkpoints( preflight_participants.append( ( checkpoint_path, - table, + tables, attempted_case_ids, checkpoint, effective_state, ) ) - for checkpoint_path, table, attempted_case_ids, checkpoint, effective_state in preflight_participants: + for checkpoint_path, tables, attempted_case_ids, checkpoint, effective_state in preflight_participants: canonical_path = _checkpoint_path_in_locked_directory(checkpoint_path, locked) _require_effective_checkpoint_state_at( locked.file_descriptor, @@ -4035,7 +4454,7 @@ def _validated_transaction_checkpoints( validated_participants.append( _ValidatedCheckpoint( attestation=snapshot.attest(checkpoint_path), - table=table, + tables=tables, attempted=frozenset(attempted_case_ids), document=checkpoint, ) @@ -4067,7 +4486,8 @@ def _validate_transaction_table_ownership( allowed_staging_by_table = {Path(str(perf_file)).stem: output_root / str(perf_file) for perf_file in PerfFile} checkpoints_by_table: dict[str, list[_ValidatedCheckpoint]] = {} for checkpoint in checkpoints: - checkpoints_by_table.setdefault(checkpoint.table, []).append(checkpoint) + for table in checkpoint.tables: + checkpoints_by_table.setdefault(table, []).append(checkpoint) if any(table not in allowed_staging_by_table for table in checkpoints_by_table): raise RuntimeError(f"Invalid table in collector sidecar transaction {journal_path}") expected_staging_paths = {allowed_staging_by_table[table] for table in checkpoints_by_table} @@ -5385,7 +5805,8 @@ def _commit_collector_provenance_transaction( journal_attestation=journal_attestation, locked_output_root=locked_output_root, ) - logger.info(f"Wrote collector provenance sidecar: {meta_path}") + if logger is not None: + logger.info(f"Wrote collector provenance sidecar: {meta_path}") for checkpoint in participants: _close_checkpoint_attempts( checkpoint.path, @@ -5478,15 +5899,16 @@ def _write_collector_provenance( module_by_table: dict[str, str] = {} staging_by_table: dict[str, Path] = {} for collection in collections: - table = Path(str(collection["perf_filename"])).stem full_name = f"{collection['name']}.{collection['type']}" - ops_by_table.setdefault(table, []).append(full_name) collection_by_full_name[full_name] = collection - module_by_table.setdefault(table, collection["module"]) - staging_path = Path(str(collection["perf_filename"])) - if not staging_path.is_absolute(): - staging_path = output_root / staging_path - staging_by_table.setdefault(table, staging_path) + for perf_filename in _collection_perf_filenames(collection): + table = Path(perf_filename).stem + ops_by_table.setdefault(table, []).append(full_name) + module_by_table.setdefault(table, collection["module"]) + staging_path = Path(perf_filename) + if not staging_path.is_absolute(): + staging_path = output_root / staging_path + staging_by_table.setdefault(table, staging_path) module_failure_names = {e["module"] for e in run_errors if e.get("error_type") == "ModuleCollectionFailure"} checkpoint_root = _checkpoint_backend_root(checkpoint_dir, backend) @@ -5506,7 +5928,10 @@ def _write_collector_provenance( full_names = ops_by_table.get(table) module = module_by_table.get(table) if not full_names or module is None: - logger.warning(f"collection_meta: {table} has no registry mapping this run; skipping its provenance entry") + if logger is not None: + logger.warning( + f"collection_meta: {table} has no registry mapping this run; skipping its provenance entry" + ) continue if parquet_path not in finalization_info: @@ -5544,7 +5969,7 @@ def _write_collector_provenance( ) if resume_tracker is None: continue - checkpoint_plan = _producer_checkpoint_plan(resume_tracker, table) + checkpoint_plan = _producer_checkpoint_plan(resume_tracker, _collection_tables(collection)) else: resume_tracker = _resume_tracker_for_collection( collection, @@ -5556,7 +5981,7 @@ def _write_collector_provenance( checkpoint_plan = producer_plan.get(resume_tracker._path) if ( checkpoint_plan is None - or checkpoint_plan.table != table + or table not in checkpoint_plan.tables or checkpoint_plan.identity_dict() != resume_tracker._metadata ): raise RuntimeError( @@ -5902,6 +6327,25 @@ def main(): help="Run only cases whose string form contains SUBSTR (repeatable, OR semantics). " "Ephemeral healing filter — never persisted to YAML.", ) + parser.add_argument( + "--cases-from", + type=str, + default=None, + metavar="FILE", + help="Failure replay: run exactly the cases listed in FILE — an errors_*.json[.gz] or " + "collection_summary_*.json[.gz] from an earlier run, a JSON list of case strings, or one case per " + "line — matched by exact case string. Writes collector_replay.json beside the results. A replay " + "is a diagnosis: it needs an output directory without collector results and keeps its rows as CSV " + "staging (no parquet, no collection_meta.yaml).", + ) + parser.add_argument( + "--replay-control", + type=int, + default=0, + metavar="N", + help="With --cases-from: also run N random cases per op that were NOT requested (seed 42) as a " + "control group for before/after comparison.", + ) parser.add_argument( "--plan-only", action="store_true", @@ -5972,9 +6416,16 @@ def main(): ops = [op for op in planned_ops if op in requested_ops] missing_ops = requested_ops - set(ops) if missing_ops: - parser.error( - "Requested ops are not present in the collector v2 case plan: " + ", ".join(sorted(missing_ops)) + # Not a usage error: the op is registered but the model plan keeps it + # registry-only (or the caller's shard plan is stale). A dedicated exit + # code lets a scheduler mark the shard terminal instead of retrying an + # argparse 2 until max_attempts (2026-10-11, sm100 shard plans). + print( + f"{parser.prog}: error: Requested ops are not present in the collector v2 case plan: " + + ", ".join(sorted(missing_ops)), + file=sys.stderr, ) + sys.exit(EXIT_CODE_OPS_NOT_IN_PLAN) if (args.model_path or args.model_architecture) and not case_plan.model_cases_paths: logger_message = ( @@ -5997,8 +6448,16 @@ def main(): else: os.environ.pop("COLLECTOR_MODEL_PATH", None) + # helper.get_sm_version lets AIS_SM stand in for a missing GPU when enumerating cases; + # under a collection it would dispatch and label every case for that SM instead. + forced_sm = os.environ.get("AIS_SM") or os.environ.get("AIC_SM") + if forced_sm: + parser.error(f"AIS_SM/AIC_SM={forced_sm} is a case-inventory override; unset it to collect") + # Setup logging - debug flag is handled inside setup_logging - if logger is None: + global _LOGGER_CONFIGURED + if not _LOGGER_CONFIGURED: + _LOGGER_CONFIGURED = True if args.model_cases_full: log_scope = ["model_cases_full"] else: @@ -6112,6 +6571,9 @@ def was_touched_by_run(path: Path) -> bool: resolved = path.resolve() return resolved not in existing_perf_outputs or path.stat().st_mtime_ns != existing_perf_outputs[resolved] + replay = _load_replay_spec(args, logger) + _require_replay_output_root_without_results(replay, output_root, existing_perf_outputs) + # Use profiling context manager with ProfilerContext(args.backend, enabled=args.profile): collect_backend = {"trtllm": collect_trtllm, "sglang": collect_sglang, "vllm": collect_vllm}[args.backend] @@ -6125,10 +6587,17 @@ def was_touched_by_run(path: Path) -> bool: case_plan=case_plan, sm_version=sm_version, case_filters=args.case_filters, + replay=replay, ) + _finish_replay(replay, output_root, logger) converted: list[Path] = [] - if args.keep_csv: + if replay is not None: + logger.info( + f"Replay run: keeping collector CSV staging files; a replay never finalizes into parquet or " + f"collection_meta.yaml ({ReplaySpec.MARKER})" + ) + elif args.keep_csv: logger.info("Keeping collector CSV staging files because --keep-csv was passed") else: touched_perf_outputs = {path for path in find_perf_csv_outputs(output_root) if was_touched_by_run(path)} diff --git a/python/aisimulate/collector/evidence_exceptions.yaml b/python/aisimulate/collector/evidence_exceptions.yaml index 578d441a9e..6efd7f3090 100644 --- a/python/aisimulate/collector/evidence_exceptions.yaml +++ b/python/aisimulate/collector/evidence_exceptions.yaml @@ -4,13 +4,14 @@ # Two registries in one file: # evidence_exceptions — waivers against resolver-emitted evidence # requirements (design §9 "Exceptions"). -# standards_exceptions — owner-approved, expiring exceptions to collector -# RULES (mechanism defined in -# .claude/rules/collector/layer_permissions.md +# standards_exceptions — owner-approved exceptions to collector RULES, +# re-verified on every pin bump (mechanism defined +# in .claude/rules/collector/layer_permissions.md # §"Governed exceptions"). # Both are human-reviewed against `evidence_check.py` output until the # enforcing CI gate lands; every entry carries a self-contained -# machine-readable scope, an approver, and an expiry. +# machine-readable scope and an approver. Evidence waivers carry an expiry; +# standards exceptions carry the release they were last verified on. schema_version: 2 @@ -80,6 +81,58 @@ evidence_exceptions: standards_exceptions: +# --- compute_scale (scale_matrix_perf) is an eager-only standalone kernel --- +# Owner decision (tianhaox, 2026-09-25): keep the compute_scale rows as an +# APPROXIMATION. Under vLLM 0.29's default execution mode (torch.compile with +# the norm_quant / act_quant fusion passes) the per-tensor fp8 activation +# quant never runs as its own kernel — it is folded into Inductor kernels +# (triton_*_fused__to_copy_clamp_cutlass_scaled_mm_* on Llama-3.1-70B-FP8, +# framework-mode probe 2026-09-24), so the op-level identity gate has no +# serving instance to compare against and the SDK bills a kernel serving +# does not launch separately. Direction: conservative (over-charges the +# quant boundary by a standalone kernel's launch + one extra pass over the +# activation). The rows stay collected as-is; consumers see the deviation +# through this entry and the opharness finding +# full_reprobe_framework_mode_2026_09_24. +- rule: serving-dispatch parity (layer_permissions.md) — compute_scale standalone quant kernel under torch.compile + scope: + tables: [scale_matrix_perf] + backend: vllm + versions: all + systems: [l40s, h100_sxm, h200_sxm] # SM89 too: op-probe-2026-10 findings.yaml, sm89 vllm-0.30.0 entry + retirement: "model the fused quant (fold compute_scale into the gemm row when the deployment compiles) or drop the standalone row for compiled deployments" + approved_by: tianhaox + verified_on: "0.30.0" # opharness compute_scale gate, 2026-09-27; compute_scale is pinned to 0.30.0 + +# --- fp8_static GEMM is DERIVED, and on CC>=100 serving runs a different kernel --- +# Owner decision (tianhaox, 2026-09-30): keep fp8_static as an APPROXIMATION, +# not a collected row — models needing it (ModelOpt static per-tensor FP8 +# checkpoints) will stay few. The SDK derives fp8_static from the dynamic +# per-token `fp8` gemm row plus compute_scale/scale_matrix +# (perf_database._gemm_key_names), which assumes both activation schemes run +# the SAME GEMM kernel. vLLM 0.30.0 breaks that on CC>=100: the static +# per-tensor scheme selects FlashInferFP8ScaledMMLinearKernel +# (model_executor/kernels/linear/__init__.py:421-430 priority; scaled_mm/ +# flashinfer.py:50-63 is_supported CC>=100 + per-tensor act/weight) while the +# dynamic scheme falls to CutlassFP8ScaledMMLinearKernel. Observed on the SM120 +# probe box: serving Llama-3.1-70B-Instruct-FP8 logs "Selected +# FlashInferFP8ScaledMMLinearKernel" and runs cutlass +# MainloopSm120TmaWarpSpecializedBlockwiseScaling + cuBLAS sm89_xmma at M=1; +# the `fp8` row runs the per-tensor Cutlass mainloop (opharness finding +# sm120_collector_gemm_fp8_static_row; gate gemm_fp8_Llama-3.1-70B-FP8 is an +# explained deviation, not a gate). Direction: unknown sign (different kernel, +# not a superset) — treat fp8_static predictions on these systems as +# approximate. Below CC 100 both schemes select Cutlass and the derivation holds. +- rule: serving-dispatch parity (layer_permissions.md) — fp8_static derived from the dynamic fp8 gemm row while CC>=100 serving selects a different scaled-mm kernel + scope: + tables: [gemm_perf] + backend: vllm + versions: all + systems: [rtx_pro_6000_server, b200_sxm, b300_sxm, gb200, gb300] # every CC>=100 system; only rtx_pro_6000_server-class (SM120) hardware-observed so far + retirement: "collect a static per-tensor fp8 gemm row (Fp8Config activation_scheme=static) on CC>=100 systems and let perf_database prefer the measured fp8_static row, or drop fp8_static support for ModelOpt FP8 checkpoints on those systems" + approved_by: tianhaox + verified_on: "0.31.0" # observed on 0.30.0; 0.31.0 source unchanged: kernels/linear/__init__.py:422-423 priority, scaled_mm/flashinfer.py:50-62 + # --- PR #1486: Blackwell mHC fused_hc serving-dispatch parity ------------- # Owner decision (tianhaox, 2026-08-19/20): the fused_hc collection on # SM100/103 is REJECTED for now. Default rc23 Blackwell mHC predictions @@ -96,8 +149,8 @@ standards_exceptions: scope: tables: [mhc_module_perf] backend: trtllm - versions: ["1.3.0rc23"] + versions: all systems: [b200_sxm, b300_sxm, gb200, gb300] retirement: "https://github.com/ai-dynamo/aiconfigurator/issues/1580 — collect fused_hc on SM100/103 + SDK consumer switch" approved_by: tianhaox - expires: 2026-10-31 + verified_on: "1.3.0rc29" # observed on rc23; rc29 serving still defaults to fused_hc (models/modeling_deepseekv4.py:292-297) while the collector times pre/post_mapping diff --git a/python/aisimulate/collector/framework_manifest.yaml b/python/aisimulate/collector/framework_manifest.yaml index 6f9afdad2a..b822d1b4c2 100644 --- a/python/aisimulate/collector/framework_manifest.yaml +++ b/python/aisimulate/collector/framework_manifest.yaml @@ -14,27 +14,18 @@ schema_version: 2 frameworks: sglang: source_repo: "https://github.com/sgl-project/sglang.git" + # 2026-10-01 (H20/sm90 upgrade audit, benchmarks/evidence/collector/op-probe-2026-10/findings.yaml + # hopper_sglang_collector_port_0514_to_0521_2026_10_01): default pin moved + # 0.5.14 -> 0.5.21 (newest release tag). Both tags resolve to the same + # multi-arch index digest, as the 0.5.17 model pins below do. The former + # kda (kimi-k3 branch build, 0.5.16) and msa (official v0.5.16) family pins + # are gone: 0.5.21 ships models/kimi_linear.py + the KDA Triton lanes and + # models/minimax_m3.py, and both op families were smoked/gated on it. default: - version: "0.5.14" + version: "0.5.21" images: - default: "lmsysorg/sglang:v0.5.14@sha256:5027e95bf6ec536856b1b52a91d1f35ff5c564ab83e8a94758a169ff09bb8df3" - cu130: "lmsysorg/sglang:v0.5.14-cu130@sha256:5027e95bf6ec536856b1b52a91d1f35ff5c564ab83e8a94758a169ff09bb8df3" - families: - # Kimi-K3 KDA kernels exist only on the sglang kimi-k3 branch - # (https://github.com/sgl-project/sglang/tree/kimi-k3), whose build - # reports 0.5.16. Pin the branch image for the kda family only. - kda: - version: "0.5.16" - images: - default: "lmsysorg/sglang@sha256:4b8a7542e00a4d82801d1276108438b024137c9a07f33215d03b04003e500f2b" - # MiniMax-M3 (minimax_m3.py + minimax_sparse_ops) landed in the official - # v0.5.16 release; the stock 0.5.14 default has no M3. Pin the official - # tag for the msa family only (NOT the kimi-k3 branch build above, which - # is a different container). - msa: - version: "0.5.16" - images: - default: "lmsysorg/sglang:v0.5.16@sha256:7b6a35df9839fd593a94a1eaee82d7777f472225d9f3ad1f8a2e0cb2bd1785d0" + default: "lmsysorg/sglang:v0.5.21@sha256:b1259f3ea3275f66237c498ea388919729018bc9f01c3d638391e06e2cf3f469" + cu130: "lmsysorg/sglang:v0.5.21-cu130@sha256:b1259f3ea3275f66237c498ea388919729018bc9f01c3d638391e06e2cf3f469" # Model-scoped collector runtime pins: when the resolved run model_path # exactly equals one of these keys, the pin overrides BOTH the family # pins above and the framework default for every op in the run — the @@ -48,12 +39,13 @@ frameworks: # one model's entire run and leaves every other model's resolution (and # every other op's family/default resolution) unchanged. # - # SGLang v0.5.17 registers Qwen3_5MoeForCausalLM / - # Qwen3_5ForCausalLM; v0.5.14 predates that model integration and cannot - # collect these cases. Both tags below resolve to the same multi-arch - # index digest. These are collector-runtime declarations, not claims of - # whole-checkpoint serving validation (in particular for the RadixArk - # NVFP4 collection fixture). + # Qwen3.8-2.4T model pins (0.5.17) folded into the 0.5.21 default + # 2026-10-30 per the pin-consolidation rule (pins below the target version + # upgrade with it): the 0.5.21 tables already carry the Qwen3.8 MoE and GDN + # shapes from the unfiltered campaign, and the op-probe serving matrix boots + # Qwen3.8-2.4T-A95B on SM90/SM100 and -FP8 on SM90 under 0.5.21 + # (benchmarks/evidence/collector/op-probe-2026-10). The 0.5.17 tables stay + # under their version key. # # Scope: the override is selected only when collect.py resolves a single # exact model_path (normally through --model-path, or through an @@ -61,126 +53,74 @@ frameworks: # has no single model identity and therefore uses the framework/family # runtime; run these model-pinned campaigns separately at their declared # runtime rather than mixing them into an unfiltered invocation. - models: - "Qwen/Qwen3.8-2.4T-A95B": - version: "0.5.17" - images: - default: "lmsysorg/sglang:v0.5.17@sha256:16aba8925507e631e1dc1e23d95d026533602591775f6a8db68b74ee99746155" - cu130: "lmsysorg/sglang:v0.5.17-cu130@sha256:16aba8925507e631e1dc1e23d95d026533602591775f6a8db68b74ee99746155" - "Qwen/Qwen3.8-2.4T-A95B-FP8": - version: "0.5.17" - images: - default: "lmsysorg/sglang:v0.5.17@sha256:16aba8925507e631e1dc1e23d95d026533602591775f6a8db68b74ee99746155" - cu130: "lmsysorg/sglang:v0.5.17-cu130@sha256:16aba8925507e631e1dc1e23d95d026533602591775f6a8db68b74ee99746155" - "RadixArk/Qwen3.8-2.4T-A95B-NVFP4": - version: "0.5.17" - images: - default: "lmsysorg/sglang:v0.5.17@sha256:16aba8925507e631e1dc1e23d95d026533602591775f6a8db68b74ee99746155" - cu130: "lmsysorg/sglang:v0.5.17-cu130@sha256:16aba8925507e631e1dc1e23d95d026533602591775f6a8db68b74ee99746155" - # GLM-5.3-Flash (Glm5NextForConditionalGeneration) standalone per-op - # collection. SGLang first ships python/sglang/srt/models/glm5_next.py - # (with dsa_indexer_kpool.py and communicator_mhc.py) in the official - # v0.5.20 release, tag commit 94602c9c2b7cbdb8efd5c52802dac6a1c180089e; - # the 0.5.14 default has no glm5_next model. Docker Hub reported the - # multi-arch index digest below for both v0.5.20 and v0.5.20-cu130 on - # 2026-09-30; its linux/arm64 child (the GB300 platform image, imported - # for collection as b0d8718a...-auth3.sqsh) is - # sha256:b0d8718a4424bb22e448e04407ab3ce5f7399a4c5fc702d6fbe36c3772ec8862. - # Scope/precedence as documented for the Qwen3.8 pins above: the pin - # applies only when collect.py resolves exactly one of these model - # paths; every other model resolves exactly as before. - "zai-org/GLM-5.3-Flash": &glm53flash_sglang_runtime - version: "0.5.20" - source_commit: "94602c9c2b7cbdb8efd5c52802dac6a1c180089e" - images: - default: "lmsysorg/sglang:v0.5.20@sha256:06e4f2ed21afde4ff513cda65070124e727ba23ccaeff7712b8c40e1097d611f" - cu130: "lmsysorg/sglang:v0.5.20-cu130@sha256:06e4f2ed21afde4ff513cda65070124e727ba23ccaeff7712b8c40e1097d611f" - "nvidia/GLM-5.3-Flash-NVFP4": *glm53flash_sglang_runtime + # GLM-5.3-Flash SGLang 0.5.20 model pin retired 2026-10-30 (owner decision: + # one release per collector file; 0.5.20 is below the 0.5.21 target). The + # 0.5.20 tables stay under their version key; re-collect them from + # main at c540ff90a (before #649). Moving the lane to 0.5.21 needs the GLM + # attention runner, KDA, mHC and GEMM ports listed in the 2026-10-30 + # source audit, then a re-collection. + models: {} trtllm: source_repo: "https://github.com/NVIDIA/TensorRT-LLM.git" + # 2026-10-02 (H20/sm90 upgrade audit, benchmarks/evidence/collector/op-probe-2026-10/findings.yaml + # hopper_trtllm_collector_pin_rc20_to_rc29_2026_10_02): default pin moved + # 1.3.0rc20 -> 1.3.0rc29 (newest release container). The rc23 moe / + # sparse_attention / msa family pins are retired: every trtllm registry op + # smokes on rc29 (21/23 clean; moe's first plan case and mamba2's largest + # cells fail by their own classified raises) and all 17 path gates align on + # it, DSV4 included. Retired-pin history, kept for the data it explains: + # * rc23 moe (2026-08-08): FP4 w4a8_mxfp4_mxfp8 rows show reproducible + # non-monotonic num_tokens segments on Blackwell (median ratio vs rc20 + # 1.003-1.013; per-bucket tactic/tile-wave quantization, not a campaign + # artifact; GitLab aic-auto-collector pipelines 61692150-61692160, + # 61771569-61838704). Existing moe shapes still reuse rc20 data through + # the declared reuse.yaml channel; the rc23 evidence exception in + # evidence_exceptions.yaml keeps its own scope. + # * rc23 sparse_attention: DSV4 attention needed trtllm>=1.3.0rc21 and the + # DSA modules travelled with it (SM90 smoke at bring-up). + # * rc23 msa: the true-sparse fmha_sm100 backend (SM100/103) landed in rc23; + # SM90 msa data stayed at rc20 (Triton reference path, identical-or-worse + # measured at rc23). default: - version: "1.3.0rc20" + version: "1.3.0rc29" images: - default: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20@sha256:1532b38814b3faf2affdb5ef01ca91468685d314ffb7e8926a0567595355ed88" - families: - # trtllm FP4 (w4a8_mxfp4_mxfp8) MoE rows show non-monotonic - # num_tokens segments (4096->8192, 12288->16384 dips) on all four - # Blackwell systems. An rc23 re-collection reproduced the dips - # (2026-08-08; exact consumer-key join vs rc20: median ratio - # 1.003-1.013, median |per-row delta| 2.4-2.9%, ~25% of rows >5% — - # benchmark run-to-run spread, not a systematic shift), so the dips - # are REAL, reproducible kernel behavior (per-bucket tactic/tile-wave - # quantization), not a campaign artifact. Evidence: GitLab - # aic-auto-collector (project 205688) pipelines 61692150-61692160 - # (moe re-collection, 2026-08-08); DSV4 attention campaigns ran as - # pipelines 61771569-61838704 on feat/sbatch-partition-override. The DSV4 moe shapes stay collected on this - # runtime for version coherence with the DSV4 attention family; the - # remaining moe shapes reuse rc20 via the declared reuse.yaml channel. - moe: - version: "1.3.0rc23" - images: - default: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc23@sha256:316b840a08a8174fc3f6b5716828bdfe1daaf629ee1ac2a8b7a22526d141a007" - # DeepSeek-V4 attention modules require trtllm>=1.3.0rc21 (the model - # class and sparse/deepseek_v4 backend first ship there; see - # collector/trtllm/collect_dsv4_attn.py __compat__). The catalog family - # also carries the DSA modules (dsa_*_module), so trtllm DSA collection - # moves to this runtime with it — DSA on rc23 smoke-verified on SM90 - # (H20) at bring-up; the rc20 serving citations in collect_mla_module - # are re-audited on this bump per layer_permissions.md. - sparse_attention: - version: "1.3.0rc23" - images: - default: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc23@sha256:316b840a08a8174fc3f6b5716828bdfe1daaf629ee1ac2a8b7a22526d141a007" - # MiniMax-M3: the true-sparse fmha_sm100 backend (implementation="msa", - # SM100/103-only) landed in rc23; rc20 has only the Triton reference - # path (measured identical-or-worse on SM90 at rc23, so the existing - # existing SM90 msa data stays at rc20). Pin the msa family to rc23 for the - # Blackwell collection. - msa: - version: "1.3.0rc23" - images: - default: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc23@sha256:316b840a08a8174fc3f6b5716828bdfe1daaf629ee1ac2a8b7a22526d141a007" + default: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc29@sha256:124da3f10b73afbf4242f5c68fb04b48736c61fbe662e715202371fd7f046f43" vllm: source_repo: "https://github.com/vllm-project/vllm.git" default: - version: "0.24.0" + version: "0.30.0" images: - default: "vllm/vllm-openai:v0.24.0@sha256:251eba5cc7c12fed0b75da22a9240e582b1c9e39f6fbc064f86781b963bd814f" - cu129: "vllm/vllm-openai:v0.24.0-cu129@sha256:6b005c78cd17a1f0215f31e95d183833f4dfd6167c4473241b80d1c8843b91e8" + default: "vllm/vllm-openai:v0.30.0@sha256:8a69ffad015f138d7170c4ddc429e230a3bc1c1719f67e14324749df200a4b90" + # gemm / moe / kda / mhc move to vLLM 0.31.0 (owner decision 2026-10-30): + # GLM-5.3-Flash collects these families on 0.31.0, and each collector file + # targets exactly one release, so the families upgrade instead of carrying + # 0.30.0 and 0.31.0 side by side. A family pin moves the family for every + # model. The 0.30.0 tables stay under their version key; the 0.30.0 + # collectors are `git log` away. Same stock image as the GLM-5.3-Flash + # model pin below (multi-arch index; amd64 a4a4c043..., arm64 3f7dd5b7...). families: - # Kimi-K3 KDA kernels exist on the official kimi-k3 preview image - # (CUDA 13 build with fused_kda_decode + FlashKDA; setuptools-scm - # version 0.1.dev19262+gb6bbf29dd) and upstreamed in v0.27.0 — but the - # family pin STAYS on the preview until the coordinated consumer fix - # lands (review jasonqinzhou #1559, comment 3808798129): at 0.27.0 the - # DS conv-state layout makes is_fused_kda_decode_supported false, so - # serving decode runs causal_conv1d_update + - # fused_recurrent_kda_packed_decode + explicit gated RMSNorm, while - # KimiK3Model._build_generation_ops still prices every vLLM non-spec - # decode as fused_kda_decode and drops the norm. The 0.27.0 gb300 kda - # rows collected by this PR are checked in as evidence - # (systems/data/gb300/kda/vllm/0.27.0); until the consumer routes by - # version, loading that version fills 44 fused_kda_decode rows from - # this preview sibling via the shared layer and misprices generation - # KDA. Re-pin when the consumer emits the conv-update + packed - # decode pair (and blocks preview fused fill) for 0.27.0+. - kda: - version: "0.1.dev19262" - images: - default: "vllm/vllm-openai@sha256:e90e2603b2781936651ba019804137714367c69e10a7b25a2e57b46995225616" - # mla: pin 0.27.0 so gb300/vllm MLA module rows are collected on the - # same runtime as the K3 serving lane (KDA-family sibling) instead of - # silently consuming the trtllm 1.3.0 low-fidelity fallback. Collector - # API drift 0.24.0 -> 0.27.0 is audited before the pin lands in data. - mla: - version: "0.27.0" - images: - default: &vllm_0270_image "vllm/vllm-openai:v0.27.0@sha256:07ea4e292adf3a26b05ac97114b28849cf4551a26beb1fbe7decd3842d752ed7" - mla_bmm: - version: "0.27.0" + gemm: &vllm_0310_runtime + version: "0.31.0" + source_commit: "db9527a46873454610df6dbedf79a36d6bf1a7f6" images: - default: *vllm_0270_image - # moe: NO family pin (owner decision 2026-08-02, reverting the + default: "vllm/vllm-openai:v0.31.0@sha256:c1c9f6fd5c109ba7f0546a59f5b2f15fb87f64c77782e90a27b648b42a8e67c3" + moe: *vllm_0310_runtime + kda: *vllm_0310_runtime + mhc: *vllm_0310_runtime + # kda family pin RETIRED (owner decision 2026-09-28): the family ran on the + # kimi-k3 preview image (0.1.dev19262) while the consumer's fused-decode + # pricing did not match 0.27's packed-decode serving path (review #1559). + # On 0.30.0 the opharness framework-mode probe of Kimi-K3 runs the fused + # decode kernel (kda_decode_fusion_many_heads_kernel) and FlashKDA prefill + # (opharness findings vllm_0300_upgrade_2026_09_27), and the kda collector + # follows 0.30's is_flashkda_supported contract — so kda resolves to the + # default 0.30.0 like every other vllm op. + # mla / mla_bmm family pins (0.27.0) were folded into the 0.29.0 + # default 2026-09-20 per the pin-consolidation rule: pins BELOW the + # target version upgrade with it; the kda preview pin (an odd dev-build + # pin) was revisited at the 0.30.0 move and retired, see above. + # moe: the 0.31.0 family pin above is a stock release (2026-10-30). + # History: NO family pin (owner decision 2026-08-02, reverting the # 2026-08-01 family-wide preview pin — a temporary preview build must # not own a family default; a family pin here hijacks image selection # for every model). ALL vllm moe collection — including K3 — runs on @@ -205,34 +145,19 @@ frameworks: # frameworks.sglang.models above: an exact run model_path overrides both # family pins and the framework default for every op in that run. # - # Qwen3.8-Max needs vLLM v0.27.1. The vLLM model registry first maps the - # bare Qwen3_5MoeForCausalLM architecture string at v0.27.1 - # (model_executor/models/registry.py:202 at tag v0.27.1, commit - # 6e448d0ea9bf3d88d898b65449ca6dc2aec170ac); v0.24.0 has the model class - # but no standard registry route for that architecture. Docker Hub's - # registry API reported the immutable index digests below on 2026-08-21. - # v0.27.1 has default and cu129 variants, but no cu128/cu130 tags. + # Qwen3.8-2.4T model pins (0.27.1) folded into the 0.30.0 default + # 2026-10-30 per the pin-consolidation rule: v0.30.0 registers + # Qwen3_5MoeForCausalLM, the 0.30.0 tables already carry the Qwen3.8 MoE + # and GDN shapes, and the op-probe serving matrix boots -FP8 (SM90) and + # nvidia's NVFP4 checkpoint (SM90/SM100) under 0.30.0 + # (benchmarks/evidence/collector/op-probe-2026-10). + # The 0.27.1 tables stay under their version key. # # Scope: the override is selected only when collect.py resolves one exact # model_path. A full/unfiltered plan has no single model identity and uses # the normal framework/family runtime; run these campaigns separately at # their declared runtime. models: - "Qwen/Qwen3.8-2.4T-A95B": - version: "0.27.1" - images: - default: "vllm/vllm-openai:v0.27.1@sha256:0a51ea5b4ae2dc5d81890e5173f54203d2a3ae0cfffe51b8fd2afd4391bfd967" - cu129: "vllm/vllm-openai:v0.27.1-cu129@sha256:07913e94a58a4e61322c88f8d4647411b4dd394838b53abacaa32448d9d8de0a" - "Qwen/Qwen3.8-2.4T-A95B-FP8": - version: "0.27.1" - images: - default: "vllm/vllm-openai:v0.27.1@sha256:0a51ea5b4ae2dc5d81890e5173f54203d2a3ae0cfffe51b8fd2afd4391bfd967" - cu129: "vllm/vllm-openai:v0.27.1-cu129@sha256:07913e94a58a4e61322c88f8d4647411b4dd394838b53abacaa32448d9d8de0a" - "RadixArk/Qwen3.8-2.4T-A95B-NVFP4": - version: "0.27.1" - images: - default: "vllm/vllm-openai:v0.27.1@sha256:0a51ea5b4ae2dc5d81890e5173f54203d2a3ae0cfffe51b8fd2afd4391bfd967" - cu129: "vllm/vllm-openai:v0.27.1-cu129@sha256:07913e94a58a4e61322c88f8d4647411b4dd394838b53abacaa32448d9d8de0a" # GLM-5.3-Flash (Glm5NextForConditionalGeneration) standalone per-op # collection runs on the stock official v0.31.0 release (no overlay), # tag commit db9527a46873454610df6dbedf79a36d6bf1a7f6. v0.31.0 moves diff --git a/python/aisimulate/collector/fullnode.py b/python/aisimulate/collector/fullnode.py index 08762a85ca..6f0d5402e9 100644 --- a/python/aisimulate/collector/fullnode.py +++ b/python/aisimulate/collector/fullnode.py @@ -75,6 +75,7 @@ def collect_sglang_fullnode_op( get_test_cases_for_model, resume_checkpoint_cls, logger, + replay=None, ) -> list[dict]: """Run one SGLang full-node op with normal V3 checkpoint evidence.""" op = collection["type"] @@ -107,6 +108,14 @@ def collect_sglang_fullnode_op( before_count = len(cases) cases = [case for case in cases if any(fragment in str(case) for fragment in case_filters)] logger.info(f"{op}: --case-filter kept {len(cases)}/{before_count} cases") + if replay is not None: + # same selection, in the same order, as the pool ops (collect.py get_func_with_limit) + before_count = len(cases) + cases = replay.select(op, cases, random.Random(shuffle_seed)) + logger.info( + f"{op}: replay selected {len(replay.matched[op])} requested + " + f"{replay.controls[op]} control of {before_count} cases" + ) if shuffle: rng = random.Random(shuffle_seed) rng.shuffle(cases) diff --git a/python/aisimulate/collector/hash_closures.yaml b/python/aisimulate/collector/hash_closures.yaml index 53f899ad09..b678f06d41 100644 --- a/python/aisimulate/collector/hash_closures.yaml +++ b/python/aisimulate/collector/hash_closures.yaml @@ -131,8 +131,6 @@ collector.sglang.collect_kda: - collector/cases/base_ops/kda.yaml - __model_cases__ collector.sglang.collect_gemm: - # SGLang 0.5.20 runs are delegated to collect_gemm_0520.py. - - collector/sglang/collect_gemm_0520.py - collector/cases/base_ops/gemm.yaml - __model_cases__ collector.sglang.collect_mhc_module: @@ -278,13 +276,6 @@ collector.vllm.collect_mla_module: - collector/vllm/utils.py - collector/cases/base_ops/mla_module.yaml - __model_cases__ -collector.vllm.collect_mla_module_027: - - collector/vllm/utils.py - - collector/vllm/collect_mla_module.py - # 0.27.0-route shim whose get/run functions are re-exported from - # collector/vllm/collect_mla_module.py — identical case reads. - - collector/cases/base_ops/mla_module.yaml - - __model_cases__ collector.vllm.collect_moe: - collector/vllm/utils.py - collector/cases/base_ops/moe.yaml diff --git a/python/aisimulate/collector/helper.py b/python/aisimulate/collector/helper.py index 38a307db70..f503fb7e97 100644 --- a/python/aisimulate/collector/helper.py +++ b/python/aisimulate/collector/helper.py @@ -24,6 +24,7 @@ import shutil import signal import stat +import subprocess import sys import tempfile import threading @@ -40,6 +41,10 @@ # Exit codes EXIT_CODE_RESTART = 10 # Exit code to indicate restart is needed +# collect.py exits with this when --ops names an op the collector v2 case plan does not +# contain (registry-only op, or a stale shard plan). Distinct from argparse's 2 so a +# scheduler can mark the shard terminal instead of retrying it. +EXIT_CODE_OPS_NOT_IN_PLAN = 3 PERF_TRANSACTION_FILENAME = ".perf-finalization.transaction.json" SIDECAR_TRANSACTION_FILENAME = ".collection_meta.transaction.json" SIDECAR_STAGING_FILENAME = ".collection_meta.pending.yaml" @@ -72,6 +77,90 @@ class PerfLogError(RuntimeError): _NVML_INITIALIZED = False _NVML_LOCK = threading.Lock() +# Stall-watchdog heartbeat. collect.py's worker registers a sink in each worker +# process and benchmark_with_power beats once per measurement, so one task that +# measures hundreds of points (a DSA context series over the whole isl x prefix +# grid) keeps counting as progress, while a point wedged inside a kernel still +# goes silent and gets the watchdog's kill after STALL_TIMEOUT_SEC. +_WORKER_HEARTBEAT: Callable[[], None] | None = None + + +def set_worker_heartbeat(sink: Callable[[], None] | None) -> None: + """Register (or clear) the callable that reports liveness to the executor.""" + global _WORKER_HEARTBEAT + _WORKER_HEARTBEAT = sink + + +def worker_heartbeat() -> None: + """Report that the running task is still measuring; a no-op outside the executor.""" + global _WORKER_HEARTBEAT + sink = _WORKER_HEARTBEAT + if sink is None: + # collect.py imports this module as ``helper`` while the op modules import + # ``collector.helper``: two module objects, so a sink registered through + # one has to be found from the other. + for name in ("helper", "collector.helper"): + other = sys.modules.get(name) + if other is not None and getattr(other, "_WORKER_HEARTBEAT", None) is not None: + sink = _WORKER_HEARTBEAT = other._WORKER_HEARTBEAT + break + if sink is None: + return + try: + sink() + except Exception: # liveness must never fail a measurement + pass + + +def _die_with_worker_preexec(): + """A Popen preexec_fn that gives the child SIGKILL when the launching worker dies + (Linux PR_SET_PDEATHSIG), or None elsewhere. prctl is resolved here, in the parent, + so the forked child only makes the call.""" + if not sys.platform.startswith("linux"): + return None + prctl = ctypes.CDLL(None, use_errno=True).prctl + pr_set_pdeathsig = 1 + + def _preexec(): + prctl(pr_set_pdeathsig, signal.SIGKILL) + + return _preexec + + +def run_measurement_subprocess(argv, *, env, cwd, timeout=None, poll_sec: float = 60.0) -> tuple[int, bytes]: + """Run a measurement child; beat the executor's stall watchdog while it makes progress. + + Progress is the child's combined stdout/stderr growing between polls: the module + collectors print every point, and PYTHONUNBUFFERED keeps that output flowing. A child + that prints nothing for the watchdog's STALL_TIMEOUT_SEC is wedged, and the watchdog + reclaims the worker; the child dies with it, so it never keeps measuring next to the + replacement worker's child. Returns ``(returncode, output)``; raises + ``subprocess.TimeoutExpired`` (with the output so far) after ``timeout`` seconds. + """ + env = {**env, "PYTHONUNBUFFERED": "1"} + with tempfile.TemporaryFile() as out: + proc = subprocess.Popen( + argv, env=env, cwd=cwd, stdout=out, stderr=subprocess.STDOUT, preexec_fn=_die_with_worker_preexec() + ) + deadline = None if timeout is None else time.monotonic() + timeout + seen = 0 + while True: + try: + proc.wait(timeout=poll_sec) + break + except subprocess.TimeoutExpired: + if deadline is not None and time.monotonic() >= deadline: + proc.kill() + proc.wait() + out.seek(0) + raise subprocess.TimeoutExpired(argv, timeout, output=out.read()) from None + size = os.fstat(out.fileno()).st_size + if size > seen: + seen = size + worker_heartbeat() + out.seek(0) + return proc.returncode, out.read() + def _parse_bool_env(env_var: str, default: bool = False) -> bool: """ @@ -207,6 +296,97 @@ def _monitoring_loop(self): self._stop_event.wait(self.interval_s) +class PhaseClock: + """Wall-clock split of one collector case or sweep point into named phases. + + ``mark(name)`` closes the phase that began at the previous mark (or at + construction); ``absorb`` folds a nested breakdown (for example the + ``phase_seconds`` that ``benchmark_with_power`` reports) into the clock + under a prefix; ``add_to`` accumulates the phases into a per-sweep total + that ``render_phase_totals`` prints. Collectors print these lines so the + campaign log shows where a case spends its seconds; they never change + what is measured. + """ + + def __init__(self) -> None: + self.start = self.last = time.perf_counter() + self.phases: dict[str, float] = {} + + def mark(self, name: str) -> None: + now = time.perf_counter() + self.phases[name] = self.phases.get(name, 0.0) + (now - self.last) + self.last = now + + def absorb(self, phases, prefix: str = "") -> None: + for name, seconds in (phases or {}).items(): + key = f"{prefix}{name}" + self.phases[key] = self.phases.get(key, 0.0) + float(seconds) + + def total(self) -> float: + return time.perf_counter() - self.start + + def render(self) -> str: + parts = [f"{name}={seconds:.2f}s" for name, seconds in self.phases.items()] + return " ".join(parts + [f"total={self.total():.2f}s"]) + + def add_to(self, totals: dict[str, float]) -> None: + for name, seconds in self.phases.items(): + totals[name] = totals.get(name, 0.0) + seconds + totals["total"] = totals.get("total", 0.0) + self.total() + totals["points"] = totals.get("points", 0) + 1 + + +def render_phase_totals(totals: dict[str, float]) -> str: + points = int(totals.get("points", 0)) + parts = [f"{name}={seconds:.1f}s" for name, seconds in totals.items() if name != "points"] + return f"{points} point(s): " + " ".join(parts) + + +class PeriodicGarbageCollector: + """Run ``gc.collect()`` every ``every`` calls, or sooner when allocated CUDA + memory sits more than ``growth_bytes`` above the lowest level seen after a + collection. + + A full collection in a process that holds a loaded model costs 0.2-0.4 s; + running it after every case or sweep point was half of an SGLang DSV4 + context sweep and 0.17 s of every TRT-LLM module case. The collectors + release framework state (KV pools, CUDA graphs, workspace buffers) + explicitly, so the collection only has to catch reference cycles, and the + memory guard keeps a cycle-heavy sweep from growing unbounded. + """ + + def __init__(self, every: int = 16, growth_bytes: int = 2 << 30) -> None: + self.every = max(1, int(every)) + self.growth_bytes = int(growth_bytes) + self.calls = 0 + self.floor: int | None = None + + @staticmethod + def _allocated() -> int: + try: + import torch + + if torch.cuda.is_available(): + return int(torch.cuda.memory_allocated()) + except Exception: + pass + return 0 + + def maybe_collect(self) -> bool: + import gc + + self.calls += 1 + allocated = self._allocated() + if self.floor is None: + self.floor = allocated + if self.calls < self.every and allocated <= self.floor + self.growth_bytes: + return False + gc.collect() + self.calls = 0 + self.floor = min(self.floor, self._allocated()) + return True + + @contextmanager def benchmark_with_power( device, @@ -220,6 +400,7 @@ def benchmark_with_power( use_cuda_graph: bool = True, # Set False only for explicitly eager measurements. *, explicit_graph_cleanup: bool = False, + num_eager_warmups: int | None = None, ): """ Context manager that handles warmup, graph capture, timing, and power monitoring. @@ -243,6 +424,10 @@ def benchmark_with_power( when capture fails or a traceback retains it. Opted into by the vLLM MLA/DSA collector; other callers retain their existing graph lifetime and synchronization behavior. + num_eager_warmups: Eager passes before capture (default: num_warmups). Graph + replay warmups stay at num_warmups. Set lower when the + kernel already ran (a dry run) and the measurement is + graph-based; eager measurements should keep the default. Yields: dict with keys: @@ -255,6 +440,8 @@ def benchmark_with_power( import torch dev_mod = get_device_module(device) + # One beat per measurement: the executor's stall watchdog counts it as progress. + worker_heartbeat() # Auto-detect configuration from environment if not explicitly provided if measure_power is None: @@ -264,19 +451,35 @@ def benchmark_with_power( # Adaptive num_runs calculation actual_num_runs = num_runs + # Eager passes before capture. They exist so lazily initialised paths (JIT, + # autotune, plan caches) run outside the capture; a caller whose kernel + # already executed (a dry run) can ask for fewer than the replay warmups, + # since the Python-heavy eager forward costs several times a graph replay. + eager_warmups = num_warmups if num_eager_warmups is None else max(int(num_eager_warmups), 0) + # Wall time of each protocol phase, reported as results["phase_seconds"]. + phase_seconds: dict[str, float] = {} + _phase_start = time.perf_counter() + + def _close_phase(name: str) -> None: + nonlocal _phase_start + now = time.perf_counter() + phase_seconds[name] = phase_seconds.get(name, 0.0) + (now - _phase_start) + _phase_start = now + if measure_power: # Estimate single iteration time with warmup start_warmup = dev_mod.Event(enable_timing=True) end_warmup = dev_mod.Event(enable_timing=True) + estimate_runs = max(eager_warmups, 1) dev_mod.synchronize() start_warmup.record() - for _ in range(num_warmups): + for _ in range(estimate_runs): kernel_func() end_warmup.record() dev_mod.synchronize() - single_iter_time = start_warmup.elapsed_time(end_warmup) / num_warmups / 1000.0 # seconds + single_iter_time = start_warmup.elapsed_time(end_warmup) / estimate_runs / 1000.0 # seconds # Adaptive duration: use shorter duration for very fast kernels to reduce memory pressure target_duration = power_min_duration @@ -292,10 +495,12 @@ def benchmark_with_power( else: # Normal warmup dev_mod.synchronize() - for _ in range(num_warmups): + for _ in range(eager_warmups): kernel_func() dev_mod.synchronize() + _close_phase("eager_warmup") + # ═══════════════════════════════════════════════════════════════════ # Device Graph Capture with Optional Fallback (CUDA: CUDAGraph, XPU: XPUGraph) # ═══════════════════════════════════════════════════════════════════ @@ -347,19 +552,20 @@ def benchmark_with_power( dev_mod.empty_cache() use_graph = False + _close_phase("capture") capture_phase_completed = True # ═══════════════════════════════════════════════════════════════ # Warmup the ACTUAL execution path (after graph capture) # ═══════════════════════════════════════════════════════════════ - dev_mod.synchronize() - for _ in range(num_warmups): - if use_graph: + # Eager callers already ran their warmups before the capture step; + # repeating them here re-ran the same eager path (23% of an SGLang + # MSA context sweep) ahead of an identical timed loop. + if use_graph: + dev_mod.synchronize() + for _ in range(num_warmups): g.replay() - else: - # Fallback: Direct execution matching actual execution path - for _ in range(repeat_n): - kernel_func() - dev_mod.synchronize() + dev_mod.synchronize() + _close_phase("graph_warmup") # Initialize power monitor if enabled power_monitor = None @@ -396,6 +602,7 @@ def benchmark_with_power( kernel_func() end_event.record() dev_mod.synchronize() + _close_phase("timed") # Check for throttling throttled = False @@ -428,6 +635,7 @@ def benchmark_with_power( "throttled": throttled, "num_runs_executed": actual_num_runs, "used_cuda_graph": use_graph, # NEW: Inform caller which path was used + "phase_seconds": phase_seconds, } finally: # Preserve default cleanup after timing starts. Explicit cleanup also @@ -440,6 +648,7 @@ def benchmark_with_power( if not explicit_graph_cleanup: dev_mod.synchronize() dev_mod.empty_cache() + _close_phase("cleanup") @contextmanager @@ -474,6 +683,411 @@ def power_monitoring_only(device, measure_power: bool | None = None): pass +# Set to 0/false to time FlashInfer-backed kernels on FlashInfer's untuned default +# tactic (the pre-2026-10-27 behavior). Default on. +FLASHINFER_AUTOTUNE_ENV = "COLLECTOR_FLASHINFER_AUTOTUNE" +FLASHINFER_AUTOTUNE_BUDGET_ENV = "COLLECTOR_FLASHINFER_AUTOTUNE_BUDGET" +FLASHINFER_AUTOTUNE_BUDGET_FLASHINFER = "flashinfer" + + +@dataclass(frozen=True) +class FlashInferTuneBudget: + """Per-tactic profiling budget applied to FlashInfer's autotuner during the tuning pass. + + FlashInfer profiles every candidate tactic with ``warmup`` untimed launches, ``repeat`` + timed launches and a ``stream_delay_us`` delay kernel in front of the timed block (its + defaults: 3 / 10 / 5000, enough to rank tactics for a server that tunes once per + process). The collector tunes once per (shape, EP, num_tokens), so that budget is the + dominant cost of a FlashInfer-backed sweep (the 2026-10-27 b200 probe: +3.6 s per + row, 28x the untuned collection time). The collector only needs the *ranking* -- the + final measurement is still ``benchmark_with_power``'s own warm-ups and runs -- and the + tactics that matter differ by 2x, so a smaller budget ranks them just as well while + near-ties stay inside collector noise. + """ + + warmup: int = 1 + repeat: int = 3 + stream_delay_us: int = 1000 + + def __str__(self) -> str: + return f"warmup={self.warmup} repeat={self.repeat} delay_us={self.stream_delay_us}" + + +def flashinfer_autotune_enabled() -> bool: + return _parse_bool_env(FLASHINFER_AUTOTUNE_ENV, default=True) + + +def flashinfer_autotune_budget() -> FlashInferTuneBudget | None: + """``COLLECTOR_FLASHINFER_AUTOTUNE_BUDGET``: ``"warmup,repeat,delay_us"`` (default + ``"1,3,1000"``), or ``"flashinfer"`` to leave FlashInfer's own 3/10/5000 untouched + (returns ``None``).""" + raw = os.environ.get(FLASHINFER_AUTOTUNE_BUDGET_ENV, "").strip() + if not raw: + return FlashInferTuneBudget() + if raw.lower() == FLASHINFER_AUTOTUNE_BUDGET_FLASHINFER: + return None + parts = [part.strip() for part in raw.split(",")] + try: + warmup, repeat, delay_us = (int(part) for part in parts) + except ValueError: + raise ValueError( + f"{FLASHINFER_AUTOTUNE_BUDGET_ENV}={raw!r}: expected 'warmup,repeat,delay_us' " + f"(three integers) or '{FLASHINFER_AUTOTUNE_BUDGET_FLASHINFER}'" + ) from None + if warmup < 0 or repeat < 1 or delay_us < 0: + raise ValueError(f"{FLASHINFER_AUTOTUNE_BUDGET_ENV}={raw!r}: warmup and delay_us must be >= 0, repeat >= 1") + return FlashInferTuneBudget(warmup, repeat, delay_us) + + +def cuda_context_error(device=None) -> str | None: + """Probe whether the CUDA context of ``device`` (default: current) is still usable. + + A sticky CUDA error (illegal memory access, launch failure) poisons the whole + context: every later kernel on this process fails, so a case that merely runs + next on the worker is reported as failed for a crash it did not cause. The probe + synchronises the device and launches one trivial kernel; it returns the error + text when either raises, ``None`` when the context is healthy or torch is absent. + """ + if str(device).startswith("xpu"): + return None # XPU workers run the same failure path; their context is not CUDA's + try: + import torch + except ImportError: + return None + try: + if device is None: + device = torch.cuda.current_device() + torch.cuda.synchronize(device) + torch.ones(1, device=device).sum().item() + except Exception as exc: # any failure here means the context is gone + return f"{type(exc).__name__}: {exc}" + return None + + +def _flashinfer_failed_tactics(tuner) -> str: + stats = getattr(tuner, "stats", None) + failed = getattr(stats, "failed_tactics", None) + if not failed: + return "no failed tactics recorded" + try: + parts = [] + for op, tactics in failed.items(): + names = sorted(str(t) for t in tactics) + shown = ", ".join(names[:4]) + (", ..." if len(names) > 4 else "") + parts.append(f"{op}: {len(names)} failed tactics ({shown})") + return "; ".join(parts) + except Exception: # diagnostics only + return f"failed tactics: {failed!r}"[:400] + + +_FLASHINFER_TUNER_BUDGET_ATTRS = ( + ("warmup", "warmup"), + ("repeat", "repeat"), + ("stream_delay_us", "stream_delay_micro_secs"), +) + + +@contextmanager +def _flashinfer_tuner_budget(tuner, budget: FlashInferTuneBudget | None): + """Apply ``budget`` to the ``AutoTuner`` singleton's plain profiling attributes for the + duration of the block, restoring whatever was there afterwards.""" + if budget is None: + yield + return + saved = {} + for field_name, attr in _FLASHINFER_TUNER_BUDGET_ATTRS: + if hasattr(tuner, attr): + saved[attr] = getattr(tuner, attr) + setattr(tuner, attr, getattr(budget, field_name)) + try: + yield + finally: + for attr, value in saved.items(): + setattr(tuner, attr, value) + + +_FLASHINFER_TUNER_DIAG_MARKERS = ( + # MoERunner.get_valid_tactics: C++ enumeration error swallowed, [] returned (DEBUG). + "failed to get valid tactics", + # Per-tactic profiling failure: "Failed when profiling ... Error + # occurred: " (DEBUG) and the "Skipped N unsupported tactic(s)" summary (INFO). + "error occurred", + "unsupported tactic", + # Nothing profiled at all: the op is served by the fallback tactic. + "no valid tactic", + "fallback", + "oom", +) + + +def _flashinfer_loggers() -> list[logging.Logger]: + """The loggers FlashInfer's autotuner writes to. ``flashinfer.jit.core.logger`` is a + ``FlashInferJITLogger`` built directly (not via ``logging.getLogger``), so it has to be + imported; the stdlib name is included for builds that register it normally.""" + loggers: list[logging.Logger] = [] + try: + from flashinfer.jit.core import logger as fi_logger + + loggers.append(fi_logger) + except Exception: # optional dependency: any import failure means "absent" + pass + std = logging.getLogger("flashinfer") + if all(lg is not std for lg in loggers): + loggers.append(std) + return loggers + + +class _LogMessageCapture(logging.Handler): + def __init__(self): + super().__init__(logging.DEBUG) + self.messages: list[str] = [] + + def emit(self, record): + try: + self.messages.append(record.getMessage()) + except Exception: # diagnostics must never break the tuning pass + pass + + +@contextmanager +def _flashinfer_tuner_log_capture(): + """Collect FlashInfer's autotuner log at DEBUG for the block without changing what it + prints: the loggers' pre-existing handlers are pinned at the old effective threshold + while the logger level is lowered, and everything is restored afterwards.""" + handler = _LogMessageCapture() + saved: list[tuple[logging.Logger, int, list[tuple[logging.Handler, int]]]] = [] + for lg in _flashinfer_loggers(): + effective = lg.getEffectiveLevel() + pinned = [] + for existing in lg.handlers: + pinned.append((existing, existing.level)) + existing.setLevel(max(existing.level, effective)) + saved.append((lg, lg.level, pinned)) + lg.addHandler(handler) + lg.setLevel(logging.DEBUG) + try: + yield handler.messages + finally: + for lg, level, pinned in saved: + lg.removeHandler(handler) + lg.setLevel(level) + for existing, old_level in pinned: + existing.setLevel(old_level) + + +def _flashinfer_tuner_diagnostics(messages: Iterable[str]) -> list[str]: + """Autotuner messages that explain a tactic *not* being profiled or chosen: the + per-runner ``get_valid_tactics`` failure (``core.py`` swallows the C++ exception and + returns ``[]`` at DEBUG), per-tactic profiling failures (every candidate runs the full + op, routing kernel included, on synthetic inputs; a failing one is marked inf at DEBUG), + empty tactic sets, fallback selection and OOM. Duplicates are collapsed.""" + out: list[str] = [] + seen: set[str] = set() + for msg in messages: + low = msg.lower() + if not any(marker in low for marker in _FLASHINFER_TUNER_DIAG_MARKERS): + continue + text = " ".join(msg.split())[:400] + if text not in seen: + seen.add(text) + out.append(text) + return out + + +def _flashinfer_failed_tactic_counts(tuner) -> dict[str, int]: + """``{op::runner: failed tactic count}`` from the tuner's statistics (empty if absent).""" + try: + failed = tuner.stats.failed_tactics + return {str(k): len(v) for k, v in failed.items()} + except Exception: # diagnostics only + return {} + + +def _flashinfer_moe_tactic_tables() -> dict[str, dict]: + """``{module: MoERunner.valid_tactics_dict}`` for the FlashInfer trtllm-gen MoE modules + already built in this process. The class-level dict caches, per instance key + ``(dtype_act, dtype_weights, fp8_quant, top_k, hidden, intermediate, local_experts, + activation, shuffled, layout, per_token, num_tokens, lora)``, the tactic list the C++ + launcher enumerated -- exactly what the autotuner had to choose from. Only modules whose + ``functools.cache`` is populated are read, so nothing is JIT-built here.""" + tables: dict[str, dict] = {} + try: + import flashinfer.fused_moe.core as fi_core + except Exception: # optional dependency + return tables + impl = getattr(fi_core, "_get_trtllm_moe_sm100_module_impl", None) + cache_info = getattr(impl, "cache_info", None) + try: + if cache_info is not None and cache_info().currsize > 0: + table = getattr(fi_core.get_trtllm_moe_sm100_module().MoERunner, "valid_tactics_dict", None) + if isinstance(table, dict): + tables["trtllm_sm100"] = table + except Exception: # diagnostics only + pass + return tables + + +def _flashinfer_new_tactic_entries(before: dict[str, set]) -> list[str]: + """Human-readable ``module key -> N tactics [...]`` lines for tactic-table keys that + appeared since ``before`` (a ``{module: set(keys)}`` snapshot).""" + lines = [] + for module, table in _flashinfer_moe_tactic_tables().items(): + for key in table: + if key in before.get(module, set()): + continue + tactics = table[key] + try: + shown = [tuple(int(x) for x in t) if isinstance(t, (list, tuple)) else t for t in list(tactics)[:8]] + count = len(tactics) + except Exception: # diagnostics only + shown, count = [repr(tactics)[:200]], -1 + more = f" (+{count - 8} more)" if count > 8 else "" + lines.append(f"{module} {key!r} -> {count} tactics {shown}{more}") + return lines + + +def _flashinfer_device_tag() -> str: + """`` dev`` for the current CUDA device, so the per-process tuning decisions of + parallel workers can be told apart in a shared log; empty without CUDA.""" + try: + import torch + + if torch.cuda.is_available(): + return f" dev{torch.cuda.current_device()}" + except Exception: # diagnostics only + pass + return "" + + +def _flashinfer_autotune_api(): + """FlashInfer's ``autotune`` context manager and ``AutoTuner`` class, or ``(None, None)``.""" + try: + from flashinfer.autotuner import AutoTuner, autotune + except ImportError: + return None, None + return autotune, AutoTuner + + +@contextmanager +def flashinfer_autotuned(num_tokens: int, tune_fn: Callable[[], None], *, label: str = ""): + """Time FlashInfer-backed kernels on the tactic FlashInfer's autotuner picks, as serving does. + + vLLM (``enable_flashinfer_autotune``, default on for Blackwell) and SGLang (warm-up + autotune, plus the extend bucket since sglang #32556) both run FlashInfer kernels on + autotuned tactics. Without this context the collector times FlashInfer's *untuned* + fallback: tile = clamp(next_pow2(num_tokens * top_k / local_num_experts), 8, 128) and + the first config of an unstable ``std::sort`` over the kernel manifest -- a value that + moved between FlashInfer 0.6.12 and 0.6.18 (flashinfer #6838) and broke every SM100 + nvfp4 MoE EP sweep whose average tokens per local expert exceeded 128 (vLLM 0.30.0 / + SGLang 0.5.21 campaign tables, 2026-10-27). + + ``tune_fn()`` runs once inside ``autotune(True, tuning_buckets=(num_tokens,))`` so + every FlashInfer op it reaches is profiled at exactly this token count (with the + caller's real routing, so the EP-local expert load is what gets tuned). The ``with`` + body then runs under ``autotune(False, ...)`` with the same bucket: lookups hit the + entries just profiled, nothing is profiled inside the body (CUDA-graph capture safe), + and an op the tuner has no entry for keeps FlashInfer's fallback exactly as before. + + The tuning pass runs under :func:`flashinfer_autotune_budget` (default 1 warm-up, 3 + timed launches, 1 ms delay per tactic instead of FlashInfer's 3 / 10 / 5 ms; see + :class:`FlashInferTuneBudget`), restored afterwards. The log line records the pass's + wall time and budget so a probe can read the tuning cost per token point. FlashInfer's + own autotuner log is captured at DEBUG during the pass and its fallback / failure + messages (an op whose ``get_valid_tactics`` raised is silently timed untuned) are + printed as a ``FlashInfer autotune WARNING`` line next to the case. + + Yields ``True`` when a tuning pass ran, ``False`` when FlashInfer is not importable or + ``COLLECTOR_FLASHINFER_AUTOTUNE`` is off (then the body runs untouched and ``tune_fn`` + is not called). Non-FlashInfer kernels inside the body are unaffected either way. + """ + autotune, auto_tuner_cls = _flashinfer_autotune_api() + if autotune is None or not flashinfer_autotune_enabled(): + yield False + return + budget = flashinfer_autotune_budget() + tuner = auto_tuner_cls.get() + cache = getattr(tuner, "profiling_cache", None) + entries_before = len(cache) if hasattr(cache, "__len__") else None + cache_keys_before = set(cache.keys()) if isinstance(cache, dict) else set() + buckets = (int(num_tokens),) + started = time.perf_counter() + tune_error: BaseException | None = None + tuner_log: list[str] = [] + failed_before = _flashinfer_failed_tactic_counts(tuner) + tactic_keys_before = {module: set(table) for module, table in _flashinfer_moe_tactic_tables().items()} + try: + with ( + _flashinfer_tuner_log_capture() as tuner_log, + _flashinfer_tuner_budget(tuner, budget), + autotune(True, tuning_buckets=buckets), + ): + tune_fn() + except Exception as exc: # re-raised below, after the context probe + tune_error = exc + elapsed = time.perf_counter() - started + # FlashInfer profiles every candidate tactic and swallows per-tactic exceptions + # (logged at DEBUG, tactic marked inf). A tactic that faults leaves a sticky CUDA + # error behind: the chosen tactic then fails with an unrelated-looking "Got CUDA + # error", and the next case on this worker inherits the dead context. Probe here so + # the failure record names the tuning pass and the worker is recycled right away. + poisoned = cuda_context_error() + if poisoned is not None: + raise RuntimeError( + f"FlashInfer autotune pass at num_tokens={num_tokens} {label} left the CUDA context " + f"in error state ({poisoned}); {_flashinfer_failed_tactics(tuner)}" + ) from tune_error + if tune_error is not None: + raise tune_error + budget_desc = str(budget) if budget is not None else "flashinfer defaults" + if entries_before is not None: + outcome = f"{len(cache) - entries_before} new tactic entries" + else: + outcome = "tuning pass done" + dev = _flashinfer_device_tag() + print( + f"FlashInfer autotune: {outcome} at num_tokens={num_tokens} {label}".rstrip() + + f" [{elapsed:.2f}s, {budget_desc}]{dev}" + ) + # An op the tuner could not enumerate tactics for is timed on the untuned fallback + # with no visible trace (b300 sglang 0.5.21 Nemotron-3-Nano nvfp4 relu2 shape, + # 2026-10-28: unchanged against the untuned table while every other shape moved). + # Surface FlashInfer's own DEBUG explanation next to the case. + # What the tuner chose: every cache entry this pass added, with its tactic. Two worker + # processes tune the same key independently (per-process cache), so a slow cell whose + # sibling distributions are fine is told apart from a slow kernel by comparing these. + if isinstance(cache, dict): + for key in cache: + if key in cache_keys_before: + continue + entry = cache[key] + tactic = entry[0] if isinstance(entry, (tuple, list)) and entry else entry + print( + f"FlashInfer autotune chosen at num_tokens={num_tokens} {label}".rstrip() + + f": {' '.join(str(key).split())[:300]} -> tactic {tactic!r}{dev}" + ) + # What the tuner had to choose from: the launcher-enumerated tactic list per instance + # key (shows the intermediate size / activation the kernel really saw and whether the + # candidate set was a single tactic, i.e. tuned == untuned by construction). + for line in _flashinfer_new_tactic_entries(tactic_keys_before): + print(f"FlashInfer autotune tactics at num_tokens={num_tokens} {label}".rstrip() + f": {line}{dev}") + diagnostics = _flashinfer_tuner_diagnostics(tuner_log) + failed_after = _flashinfer_failed_tactic_counts(tuner) + newly_failed = {k: n - failed_before.get(k, 0) for k, n in failed_after.items() if n > failed_before.get(k, 0)} + if diagnostics or newly_failed: + parts = [] + if newly_failed: + parts.append( + "tactics failed while profiling: " + ", ".join(f"{k}={n}" for k, n in sorted(newly_failed.items())) + ) + if diagnostics: + more = f" (+{len(diagnostics) - 4} more)" if len(diagnostics) > 4 else "" + parts.append(f"{len(diagnostics)} autotuner messages: " + "; ".join(diagnostics[:4]) + more) + head = f"FlashInfer autotune WARNING at num_tokens={num_tokens} {label}".rstrip() + print(f"{head}: " + " | ".join(parts) + dev) + with autotune(False, tuning_buckets=buckets): + yield True + + def setup_signal_handlers(worker_id): """Setup signal handlers to log crashes.""" import traceback as _tb @@ -499,6 +1113,18 @@ def signal_handler(signum, frame): for sig in [signal.SIGTERM, signal.SIGABRT]: signal.signal(sig, signal_handler) + # SIGUSR1 = "dump every thread's stack to stderr and keep running". The + # executor's stall watchdog sends it before killing a wedged worker, so a + # hang leaves the exact frame in the log (sglang 0.5.21 moe/bf16 on l40s + # and b300, 2026-10-06: 8 workers took their first task and never returned, + # 76 min of silence, nothing to diagnose from). + try: + import faulthandler + + faulthandler.register(signal.SIGUSR1, all_threads=True, chain=False) + except Exception: + pass + # SIGSEGV might not be catchable on all platforms try: signal.signal(signal.SIGSEGV, signal_handler) @@ -607,7 +1233,9 @@ def setup_logging(scope=["all"], debug=False, worker_id=None): # Create log directory time_stamp = datetime.now().strftime("%Y%m%d_%H%M%S") - _LOG_DIR = Path(f"{log_scope_dirname(scope)}_{time_stamp}") + # Absolute: workers and the profiler read it back from the environment, + # possibly after the working directory changed. + _LOG_DIR = Path.cwd() / f"{log_scope_dirname(scope)}_{time_stamp}" if not _LOG_DIR.is_dir(): _LOG_DIR.mkdir() @@ -678,7 +1306,19 @@ def save_error_report(errors, filename): def get_sm_version(): - """Get CUDA compute capability (SM version)""" + """Get CUDA compute capability (SM version). + + ``AIS_SM`` (``sm100`` or ``100``; legacy ``AIC_SM``) overrides the probe so + case getters can be enumerated for a target SM on a box without that GPU + (the op-probe harness's case inventory: the case set is a pure function + of (SM, framework version, case files) and must be derivable without + collecting). Real collection never sets it. + """ + forced = os.environ.get("AIS_SM") or os.environ.get("AIC_SM") + if forced: + digits = "".join(ch for ch in str(forced) if ch.isdigit()) + if digits: + return int(digits) try: import torch @@ -1911,6 +2551,7 @@ def _prepare_perf_file( snapshot_file.seek(0) table = pc_csv.read_csv(snapshot_file) _validate_moe_default_eligibility(table, parquet_path, pa=pa) + table, _duplicate_rows = _dedupe_staged_perf_rows(table, parquet_path, pa=pa, pc=pc_compute) new_rows = table.num_rows merged_existing = False merge_target = None @@ -1922,7 +2563,7 @@ def _prepare_perf_file( merge_target_was_absent = True else: old_table, merge_target = _read_attested_parquet(parquet_path, pq=pq, locked_root=locked_root) - table, merged_existing, new_rows = _merge_perf_rows(table, old_table, parquet_path, pa=pa) + table, merged_existing, new_rows = _merge_perf_rows(table, old_table, parquet_path, pa=pa, pc=pc_compute) table = _normalize_power_metrics(table, pa=pa, pc=pc_compute) temporary = perf_preparation_path(parquet_path) @@ -2566,13 +3207,71 @@ def _validate_moe_default_eligibility(table, parquet_path: Path, *, pa): raise ValueError(f"{parquet_path}: default_eligible=false requires a nonblank string kernel_source") -def _merge_perf_rows(new_table, old_table, parquet_path: Path, *, pa): +def _dedupe_staged_perf_rows(table, parquet_path: Path, *, pa, pc): + """Keep the last row per identity key within one staging file. + + The staging ``*_perf.txt`` is opened in append mode and kept across resumed + windows, and ``--resume --resume-retry-failed`` re-runs a task that failed + after logging some of its rows, so those rows are appended a second time + (2026-10-11: rtxpro msa_context 10,342 of 13,972 rows, gb200 moe/fp8_block + 2,590 pairs). The merge-with-existing path already dedupes on the identity + key; this applies the same rule to the rows of the staging file itself, + keeping the newest measurement and the original row order. Identity = every + column except the measured metrics and MoE eligibility metadata. + """ + metadata_columns = {"default_eligible"} if parquet_path.name == "moe_perf.parquet" else set() + identity = [c for c in table.column_names if c not in PERF_METRIC_COLUMNS and c not in metadata_columns] + deduped, dropped = _keep_last_row_per_identity(table, identity, pa=pa, pc=pc) + if dropped: + logging.getLogger(__name__).info( + "convert_perf_csv_to_parquet: %s: dropped %d re-appended duplicate rows " + "(%d identity keys kept the newest measurement)", + parquet_path.name, + dropped, + deduped.num_rows, + ) + return deduped, dropped + + +def _keep_last_row_per_identity(table, identity: list[str], *, pa, pc): + """Keep the last row per identity key, in the original order of the kept rows. + + pyarrow only: the collector images ship pyarrow (the harness installs it + into ``pyarrow-site``) but not pandas, and the merge path used to need + pandas — the second finalize of a shard (resume / retry-failed) died with + ``No module named 'pandas'`` (b200 vllm dsa_context_module, 2026-10-12). + Returns ``(deduped_table, dropped_rows)``. + """ + if table.num_rows < 2 or not identity: + return table, 0 + index_name = "__row_index__" + while index_name in table.column_names: + index_name = "_" + index_name + keyed = table.select(identity) + # group_by cannot hash an all-null (null-typed) column; its values carry no identity anyway + for i, field in enumerate(keyed.schema): + if pa.types.is_null(field.type): + keyed = keyed.set_column(i, field.name, keyed.column(i).cast(pa.string())) + keyed = keyed.append_column(index_name, pa.array(range(table.num_rows), type=pa.int64())) + last = keyed.group_by(identity).aggregate([(index_name, "max")]) + keep = last.column(f"{index_name}_max") + keep = pc.take(keep, pc.sort_indices(keep)) + deduped = table.take(keep) + return deduped, table.num_rows - deduped.num_rows + + +def _identity_keys(table, identity: list[str]) -> list[tuple]: + """One hashable key per row, in row order.""" + columns = [table.column(name).to_pylist() for name in identity] + return list(zip(*columns, strict=True)) if columns else [()] * table.num_rows + + +def _merge_perf_rows(new_table, old_table, parquet_path: Path, *, pa, pc): """Merge freshly-collected rows into an existing perf parquet, keeping the newest row per identity key. Returns ``(table, merged_existing, current_event_rows)``; schema incompatibility returns the new table with - ``merged_existing=False`` and its unchanged row count.""" - import pandas as pd - + ``merged_existing=False`` and its unchanged row count. pyarrow only — see + ``_keep_last_row_per_identity``.""" log = logging.getLogger(__name__) _validate_moe_default_eligibility(old_table, parquet_path, pa=pa) @@ -2595,8 +3294,8 @@ def _merge_perf_rows(new_table, old_table, parquet_path: Path, *, pa): ) # Compare (name, type) pairs for IDENTITY columns only: matching names - # with drifted types would otherwise round-trip through pandas and - # silently rewrite the parquet under a different schema. Metric columns + # with drifted types would otherwise be concatenated and silently rewrite + # the parquet under a different schema. Metric columns # are exempt from the type check — pyarrow.csv infers an all-empty # optional metric column as `null` while populated runs infer `double`, # and treating that drift as a mismatch made the merge silently OVERWRITE @@ -2604,10 +3303,21 @@ def _merge_perf_rows(new_table, old_table, parquet_path: Path, *, pa): # the finalizer then casts power metrics to double and replaces nulls with # the repository's 0.0 unavailable-measurement sentinel. Order-insensitive # (the merge realigns column order below); Arrow metadata is ignored - # (pandas round-trips change it). + # (parquets written by the pandas-based merges before 2026-10-12 carry + # pandas metadata). + # `string` and `large_string` are one logical type: pyarrow.csv infers + # `string`, but a parquet that went through one of those pandas merges + # comes back as `large_string`, so the THIRD finalize of a table (run -> + # resume -> retry-failed) used to read as a schema mismatch and OVERWRITE + # the accumulated rows (sm120 attention_generation fp8kv, 2026-10-01: + # 27,303 rows replaced by 5,067). Compare the logical string type. + def _logical(type_str: str) -> str: + return {"large_string": "string", "large_binary": "binary"}.get(type_str, type_str) + def fields(schema): return sorted( - (f.name, str(f.type)) if f.name not in PERF_METRIC_COLUMNS else (f.name, "") for f in schema + (f.name, _logical(str(f.type))) if f.name not in PERF_METRIC_COLUMNS else (f.name, "") + for f in schema ) old_fields = fields(old_table.schema) @@ -2644,30 +3354,50 @@ def fields(schema): new_table.column(f.name).cast(f.type), ) - new_df = new_table.to_pandas() - old_df = old_table.to_pandas() + # align string width variants (string vs large_string) on the new side so the + # concatenation below does not widen or fail on a logical-type-equal column + for f in old_table.schema: + if f.name in PERF_METRIC_COLUMNS: + continue + new_field = new_table.schema.field(f.name) + if new_field.type != f.type and _logical(str(new_field.type)) == _logical(str(f.type)): + new_table = new_table.set_column( + new_table.schema.get_field_index(f.name), f.name, new_table.column(f.name).cast(f.type) + ) - new_df = new_df[old_df.columns.tolist()] # align column order - identity = [c for c in old_df.columns if c not in PERF_METRIC_COLUMNS and c not in metadata_columns] + new_table = new_table.select(old_table.column_names) # align column order + identity = [c for c in old_table.column_names if c not in PERF_METRIC_COLUMNS and c not in metadata_columns] if inherit_eligibility: - annotations = old_df.drop_duplicates(subset=identity, keep="last")[identity + ["default_eligible"]] - new_df = new_df.drop(columns="default_eligible").merge(annotations, on=identity, how="left", sort=False) - flags = new_df["default_eligible"] - new_df["default_eligible"] = flags.where(flags.notna(), True).astype(bool) - new_df = new_df[old_df.columns.tolist()] - current_event_rows = len(new_df.drop_duplicates(subset=identity, keep="last")) - combined = pd.concat([old_df, new_df], ignore_index=True) - deduped = combined.drop_duplicates(subset=identity, keep="last").reset_index(drop=True) + # newest existing annotation per identity key; keys the existing table + # never annotated (or annotated null) stay eligible + annotations: dict[tuple, bool] = {} + existing_flags = old_table.column("default_eligible").to_pylist() + for key, flag in zip(_identity_keys(old_table, identity), existing_flags, strict=True): + annotations[key] = True if flag is None else bool(flag) + flags = [annotations.get(key, True) for key in _identity_keys(new_table, identity)] + new_table = new_table.set_column( + new_table.schema.get_field_index("default_eligible"), + "default_eligible", + pa.array(flags, type=pa.bool_()), + ) + current_event_rows = _keep_last_row_per_identity(new_table, identity, pa=pa, pc=pc)[0].num_rows + # Drop schema metadata on both sides: a parquet written by an earlier + # pandas-based merge carries pandas index metadata that would otherwise + # make concat_tables see two different schemas (and be re-published). + old_table = old_table.replace_schema_metadata(None) + new_table = new_table.replace_schema_metadata(None) + combined = pa.concat_tables([old_table, new_table], promote_options="default") + deduped, replaced = _keep_last_row_per_identity(combined, identity, pa=pa, pc=pc) log.info( "convert_perf_csv_to_parquet: merged %s: %d existing + %d new -> %d rows " "(%d identity keys replaced by newer measurements)", parquet_path.name, - len(old_df), - len(new_df), - len(deduped), - len(combined) - len(deduped), + old_table.num_rows, + new_table.num_rows, + deduped.num_rows, + replaced, ) - return pa.Table.from_pandas(deduped, preserve_index=False), True, current_event_rows + return deduped, True, current_event_rows def find_perf_csv_outputs(output_root: str | os.PathLike = ".", *, recursive: bool = False) -> list[Path]: diff --git a/python/aisimulate/collector/kernel_source_backends.yaml b/python/aisimulate/collector/kernel_source_backends.yaml index 885f8cdc27..b400b426d8 100644 --- a/python/aisimulate/collector/kernel_source_backends.yaml +++ b/python/aisimulate/collector/kernel_source_backends.yaml @@ -137,12 +137,13 @@ mappings: - {framework: sglang, kernel_source: sglang_torch_linear, backend: cublas, source: "PR#1342 collect_gemm.py (torch.matmul/linear)"} - {framework: sglang, kernel_source: sglang_torch_bmm, backend: cublas, source: "PR#1342 collect_mla_bmm.py (torch.bmm)"} - {framework: sglang, kernel_source: sglang_sgl_kernel_bmm_fp8, backend: cublas, source: "PR#1342 collect_mla_bmm.py (sgl_kernel bmm_fp8 = cublasLt)"} + - {framework: sglang, kernel_source: sglang_flashinfer_bmm_fp8, backend: cublas, source: "collector/sglang/collect_mla_bmm.py:22-26 (sglang>=0.5.21: fp8_utils.bmm_fp8 wraps flashinfer.bmm_fp8 as custom op flashinfer_bmm_fp8_batched, docstring 'via flashinfer's cuBLAS backend', fp8_utils.py:246-273 @0.5.21); flashinfer csrc/bmm_fp8.cu bmm_fp8_internal_cublaslt"} - {framework: sglang, kernel_source: sglang_sgl_kernel_fp8_scaled_mm, backend: cutlass, source: "PR#1342 collect_gemm.py (sgl_kernel fp8_scaled_mm)"} - {framework: sglang, kernel_source: sglang_deepgemm_gemm_nt_f8f8bf16, backend: deepgemm, source: "PR#1342 collect_gemm.py"} - - {framework: sglang, kernel_source: sglang_triton_w8a8_block_fp8, backend: triton, source: "sglang/collect_gemm_0520.py (fp8_utils.py:1094-1107 @v0.5.20 N%64/K%128 fallback)"} - - {framework: sglang, kernel_source: sglang_cutedsl_tgv_bf16, backend: cutlass, source: "sglang/collect_gemm_0520.py (CuTe DSL TGV bf16, cutedsl_bf16_gemm.py @v0.5.20)"} - - {framework: sglang, kernel_source: sglang_flashinfer_bf16_splitk, backend: flashinfer, source: "sglang/collect_gemm_0520.py (unquant.py split-K @v0.5.20)"} - - {framework: sglang, kernel_source: sglang_flashinfer_bf16_direct, backend: flashinfer, source: "sglang/collect_gemm_0520.py (unquant.py direct @v0.5.20)"} + - {framework: sglang, kernel_source: sglang_triton_w8a8_block_fp8, backend: triton, source: "sglang/collect_gemm_0520.py @c540ff90a, retired by #649; labels the shipped 0.5.20 tables (fp8_utils.py:1094-1107 @v0.5.20 N%64/K%128 fallback)"} + - {framework: sglang, kernel_source: sglang_cutedsl_tgv_bf16, backend: cutlass, source: "sglang/collect_gemm_0520.py @c540ff90a, retired by #649; labels the shipped 0.5.20 tables (CuTe DSL TGV bf16, cutedsl_bf16_gemm.py @v0.5.20)"} + - {framework: sglang, kernel_source: sglang_flashinfer_bf16_splitk, backend: flashinfer, source: "sglang/collect_gemm_0520.py @c540ff90a, retired by #649; labels the shipped 0.5.20 tables (unquant.py split-K @v0.5.20)"} + - {framework: sglang, kernel_source: sglang_flashinfer_bf16_direct, backend: flashinfer, source: "sglang/collect_gemm_0520.py @c540ff90a, retired by #649; labels the shipped 0.5.20 tables (unquant.py direct @v0.5.20)"} - {framework: sglang, kernel_source: sglang_flashinfer_cutlass_nvfp4, backend: flashinfer_cutlass, source: "PR#1342 collect_gemm.py (flashinfer mm_fp4 cutlass)"} # MoE labels name the constructed serving method + runner. - {framework: sglang, kernel_source: sglang_flashinfer_trtllm_moe, backend: flashinfer_trtllm_gen, source: "PR#1342 collect_moe.py"} @@ -152,6 +153,8 @@ mappings: # GDN sub-kernels and mHC. - {framework: sglang, kernel_source: chunk_gated_delta_rule, backend: fla_triton, source: "PR#1342 collect_gdn.py (fla triton chunk kernel)"} - {framework: sglang, kernel_source: fused_recurrent_gated_delta_rule_packed_decode, backend: fla_triton, source: "PR#1342 collect_gdn.py"} + - {framework: sglang, kernel_source: flashinfer_chunk_gated_delta_rule, backend: flashinfer, source: "collector/sglang/collect_gdn.py:366-371 (FlashInferGDNKernel.extend lane); sglang srt/layers/attention/linear/kernels/gdn_flashinfer.py:129-131,465 @0.5.21 imports flashinfer.gdn_prefill.chunk_gated_delta_rule"} + - {framework: sglang, kernel_source: fused_qkvzba_causal_conv1d_update_contiguous, backend: triton, source: "collector/sglang/collect_gdn.py:505-514 (decode fused proj + conv update); sglang kernels/ops/attention/triton_gdn_fused_proj.py:576 @0.5.21 (@triton.jit kernels), gdn_backend.py:668-760"} - {framework: sglang, kernel_source: sglang_tilelang_mhc_pre, backend: tilelang, source: "PR#1342 collect_mhc_module.py"} - {framework: sglang, kernel_source: sglang_tilelang_mhc_post, backend: tilelang, source: "PR#1342 collect_mhc_module.py"} - {framework: sglang, kernel_source: sglang_torch_mhc_pre, backend: torch, source: "PR#1342 collect_mhc_module.py (torch fallback leaf)"} @@ -220,10 +223,15 @@ mappings: # Mamba/GDN kernels. - {framework: trtllm, kernel_source: causal_conv1d_fn, backend: causal_conv1d, source: "trtllm _torch mamba modules"} + - {framework: trtllm, kernel_source: causal_conv1d_fn_token_major, backend: causal_conv1d, source: "collector/trtllm/collect_mamba2.py _run_prefill_conv_token_major: channel-last view + token-major out, the 1.3.0rc29 Mamba2Mixer default (mamba2_mixer.py:319,541-569; causal_conv1d_channellast_fwd_kernel)"} - {framework: trtllm, kernel_source: causal_conv1d_update, backend: causal_conv1d, source: "trtllm _torch mamba modules"} - {framework: trtllm, kernel_source: chunk_gated_delta_rule, backend: fla_triton, source: "trtllm _torch GDN (triton)"} - - {framework: trtllm, kernel_source: fused_recurrent_gated_delta_rule, backend: fla_triton, source: "trtllm _torch GDN (triton)"} + - {framework: trtllm, kernel_source: fused_recurrent_gated_delta_rule, backend: fla_triton, source: "trtllm _torch GDN (triton) — FLA recurrent fallback lane; serving's pure-decode entry is the fused_sigmoid_gating row below (rc23)"} + - {framework: trtllm, kernel_source: chunk_gated_delta_rule_flashinfer, backend: flashinfer, source: "collector/trtllm/collect_gdn.py: gdn_mixer._resolve_chunk_gated_delta_rule picked fla/flashinfer_chunk (FlashInfer cute-DSL GDN chunk, SM90/SM100; gdn_mixer.py:57-70 @1.3.0rc23)"} + - {framework: trtllm, kernel_source: fused_sigmoid_gating_delta_rule_update_flashinfer, backend: flashinfer, source: "collector/trtllm/collect_gdn.py:_bench_serving_decode_entry — fla/fused_sigmoid_gating_recurrent.py:452-458 dispatched to FlashInfer gdn_decode_bf16_state (kernel_cutlass_gdn_decode_*) @1.3.0rc23"} + - {framework: trtllm, kernel_source: fused_sigmoid_gating_delta_rule_update_triton, backend: fla_triton, source: "collector/trtllm/collect_gdn.py:_bench_serving_decode_entry — same entry, _can_use_flashinfer_gdn_decode False (non-SM90/100, K/V != 128, T != N) -> fused_sigmoid_gating_delta_rule_update_kernel (Triton)"} - {framework: trtllm, kernel_source: mamba_chunk_scan_combined, backend: flashinfer, source: "collector/trtllm/collect_mamba2.py:290 (mamba2 SSD chunk scan; tensorrt_llm/_torch/modules/mamba/ssd_combined.py:437-444@1.3.0rc15 flashinfer SSDCombined on SM100+ when _flashinfer_ssd_supported; Triton fallback on non-SM100 or unsupported tile sizes — no shipped row takes it: all rc15 scan rows are (chunk=128, dstate=128, headdim=64) on b200_sxm, flashinfer-eligible by construction; label split for future fallback coverage: AIC-1677)"} + - {framework: trtllm, kernel_source: selective_state_update_flashinfer, backend: flashinfer, source: "collector/trtllm/collect_mamba2.py:194-219 (_decode_state_update_kernel imports flashinfer.mamba.selective_state_update when serving selects it; mamba2_mixer.py:23 @1.3.0rc29)"} # Comm. - {framework: trtllm, kernel_source: TRTLLM, backend: custom_allreduce, source: "collector/network/collect_all_reduce.py (trtllm custom AR)"} @@ -274,7 +282,13 @@ mappings: - {framework: vllm, kernel_source: vllm_flashinfer_fidecode, backend: flashinfer, source: "PR#1344 collect_attn.py"} - {framework: vllm, kernel_source: vllm_flashinfer_trtllmprefill, backend: trtllm_mha, source: "PR#1344 collect_attn.py (flashinfer trtllm-gen prefill)"} - {framework: vllm, kernel_source: vllm_flashinfer_trtllmdecode, backend: trtllm_mha, source: "PR#1344 collect_attn.py (flashinfer trtllm-gen decode)"} - - {framework: vllm, kernel_source: vllm_flashinfer_flashinfertrtllmapidecode, backend: unverified, source: "collector/vllm/collect_attn.py:_dense_kernel_source records only the metadata class; vllm/v1/attention/backends/flashinfer.py:502-509,1539-1546@v0.25.0 uses FlashInferTrtllmAPIDecode for both XQA and TRTLLM_GEN. The label omits the kernel discriminator, so it cannot identify one backend across systems. https://github.com/vllm-project/vllm/blob/v0.25.0/vllm/v1/attention/backends/flashinfer.py"} + # vLLM 0.30.0 merged the trtllm-gen and XQA decode metadata into FlashInferTrtllmAPIDecode (flashinfer.py:608, kernel field + # FlashInferDecodeKernel xqa|trtllm-gen :568-572); collect_attn appends the kernel to the label. Both are TRT-LLM MHA kernels + # (XQA = kernel_mha in the opharness sm120 records, trtllm-gen = the SM100 fmha cubins) -> trtllm_mha. The bare class label is + # kept for rows collected before the kernel suffix (2026-10-01). + - {framework: vllm, kernel_source: vllm_flashinfer_flashinfertrtllmapidecode_xqa, backend: trtllm_mha, source: "collect_attn.py _dense_kernel_source (FlashInferTrtllmAPIDecode.kernel == XQA; sm120 fp8-KV decode, opharness serving record kernel_mha)"} + - {framework: vllm, kernel_source: vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen, backend: trtllm_mha, source: "collect_attn.py _dense_kernel_source (FlashInferTrtllmAPIDecode.kernel == TRTLLM_GEN; SM100 decode)"} + - {framework: vllm, kernel_source: vllm_flashinfer_flashinfertrtllmapidecode, backend: trtllm_mha, source: "collect_attn.py before the kernel suffix (0.30.0 rows collected 2026-09-29..10-01 on B200/RTX 5000)"} # MLA module labels: the AttentionBackendEnum member that ran. - {framework: vllm, kernel_source: FLASH_ATTN, backend: flash_attn, source: "PR#1344 collect_mla_module.py (enum-name labels)"} - {framework: vllm, kernel_source: FLASH_ATTN_MLA, backend: flash_attn, source: "PR#1344 collect_mla_module.py"} @@ -283,6 +297,9 @@ mappings: - {framework: vllm, kernel_source: FLASHMLA_SPARSE_DSV4, backend: flashmla, source: "PR#1344 collect_dsv4_attn.py"} - {framework: vllm, kernel_source: FLASHINFER_MLA, backend: flashinfer, source: "PR#1344 collect_mla_module.py"} - {framework: vllm, kernel_source: FLASHINFER_MLA_SPARSE, backend: flashinfer, source: "PR#1344 collect_mla_module.py"} + - {framework: vllm, kernel_source: FLASHINFER_MLA_SPARSE_SM120, backend: flashinfer, source: "collect_mla_module.py enum-name label; vllm v0.30.0 v1/attention/backends/registry.py:82 FlashInferMLASparseSM120Backend -> mla/flashinfer_mla_sparse_sm120.py:221-226 flashinfer_trtllm_batch_decode_with_kv_cache_mla (flashinfer.decode)"} + - {framework: vllm, kernel_source: FLASHINFER_MLA_SPARSE_DSV4, backend: flashinfer, source: "collect_dsv4_attn.py enum-name label; vllm v0.30.0 registry.py:100 DeepseekV4FlashInferMLASparseBackend -> models/deepseek_v4/nvidia/flashinfer_sparse.py:27,519 flashinfer_trtllm_batch_decode_sparse_mla_dsv4 (flashinfer.decode)"} + - {framework: vllm, kernel_source: FLASH_ATTN_MLA_SPARSE, backend: flash_attn, source: "collect_mla_module.py enum-name label; vllm v0.30.0 registry.py:119 FlashAttnMLASparseBackend -> mla/flashattn_mla_sparse.py:31,315 vllm_flash_attn flash_attn_varlen_func (fa_utils.flash_attn_supports_mla: FA3 on SM90)"} - {framework: vllm, kernel_source: TRITON_MLA, backend: triton, source: "PR#1344 collect_mla_module.py"} - {framework: vllm, kernel_source: TRTLLM_RAGGED, backend: trtllm_ragged, source: "PR#1344 collect_mla_module.py:438-444 (MLAPrefillBackendEnum.TRTLLM_RAGGED, trtllm-gen ragged prefill)"} # GEMM / quant labels: vllm kernel class names. @@ -293,6 +310,7 @@ mappings: - {framework: vllm, kernel_source: FlashInferFp8BlockScaledMMKernel, backend: flashinfer, source: "PR#1344 collect_gemm.py"} - {framework: vllm, kernel_source: FlashInferCutlassNvFp4LinearKernel, backend: flashinfer_cutlass, source: "PR#1344 collect_gemm.py"} - {framework: vllm, kernel_source: FlashInferCuteDslNvFp4LinearKernel, backend: flashinfer_cutedsl, source: "PR#1344 collect_gemm.py"} + - {framework: vllm, kernel_source: MarlinFP8ScaledMMLinearKernel, backend: marlin, source: "collect_gemm.py kernel class label; vllm v0.30.0 model_executor/kernels/linear/scaled_mm/marlin.py (apply_fp8_marlin_linear)"} - {framework: vllm, kernel_source: "torch.nn.functional.linear", backend: cublas, source: "PR#1344 collect_gemm.py (bf16)"} - {framework: vllm, kernel_source: static_scaled_fp8_quant, backend: vllm_kernel, source: "PR#1344 collect_computescale.py:73 (vllm _custom_ops)"} - {framework: vllm, kernel_source: dynamic_per_token_scaled_fp8_quant_minus_static_scaled_fp8_quant, backend: vllm_kernel, source: "PR#1344 collect_computescale.py:63 (derived overhead measure)"} @@ -305,12 +323,15 @@ mappings: - {framework: vllm, kernel_source: vllm_fp8moe_flashinfer_cutlass_flashinferexperts, backend: flashinfer_cutlass, source: "PR#1344 collect_moe.py"} - {framework: vllm, kernel_source: vllm_fp8moe_flashinfer_trtllm_trtllmfp8expertsmodular, backend: flashinfer_trtllm_gen, source: "PR#1344 collect_moe.py"} - {framework: vllm, kernel_source: vllm_fp8moe_flashinfer_trtllm_trtllmfp8expertsmonolithic, backend: flashinfer_trtllm_gen, source: "PR#1344 collect_moe.py"} + - {framework: vllm, kernel_source: vllm_fp8moe_deepgemm_deepgemmexperts, backend: deepgemm, source: "collect_moe.py:795 label = quant method + oracle backend + experts class; vllm v0.30.0 fused_moe/experts/deep_gemm_moe.py:125 DeepGemmExperts (TritonOrDeepGemmExperts leaf, triton_deep_gemm_moe.py:29)"} - {framework: vllm, kernel_source: vllm_modeloptfp8moe_triton_tritonexperts, backend: triton_fused_moe, source: "PR#1344 collect_moe.py"} - {framework: vllm, kernel_source: vllm_modeloptfp8moe_flashinfer_cutlass_flashinferexperts, backend: flashinfer_cutlass, source: "PR#1344 collect_moe.py"} - {framework: vllm, kernel_source: vllm_modeloptfp8moe_flashinfer_trtllm_trtllmfp8expertsmonolithic, backend: flashinfer_trtllm_gen, source: "PR#1344 collect_moe.py"} - {framework: vllm, kernel_source: vllm_compressedtensorsw4a4nvfp4moe_flashinfer_cutlass_flashinferexperts, backend: flashinfer_cutlass, source: "PR#1344 collect_moe.py"} - {framework: vllm, kernel_source: vllm_compressedtensorsw4a4nvfp4moe_flashinfer_trtllm_trtllmnvfp4expertsmodular, backend: flashinfer_trtllm_gen, source: "PR#1344 collect_moe.py"} - {framework: vllm, kernel_source: vllm_compressedtensorsw4a4nvfp4moe_flashinfer_trtllm_trtllmnvfp4expertsmonolithic, backend: flashinfer_trtllm_gen, source: "PR#1344 collect_moe.py"} + - {framework: vllm, kernel_source: vllm_compressedtensorsw4a4nvfp4moe_flashinfer_cutedsl_flashinfercutedslexperts, backend: flashinfer_cutedsl, source: "collect_moe.py:795; vllm v0.30.0 fused_moe/experts/flashinfer_cutedsl_moe.py:31 FlashInferCuteDSLExperts"} + - {framework: vllm, kernel_source: vllm_compressedtensorswna16moe_marlin_marlinexperts, backend: marlin, source: "collect_moe.py:795; vllm v0.30.0 fused_moe/experts/marlin_moe.py:648 MarlinExperts"} - {framework: vllm, kernel_source: vllm_modeloptnvfp4fusedmoe_flashinfer_trtllm_trtllmnvfp4expertsmonolithic, backend: flashinfer_trtllm_gen, source: "vllm/collect_moe.py GLM-5.3-Flash ModelOptNvFp4Config lane (trtllm_nvfp4_moe.py:498-520 @v0.31.0; :486-508 @v0.30.0)"} - {framework: vllm, kernel_source: vllm_compressedtensorswna16marlinmoe_marlin_marlinexperts, backend: marlin, source: "PR#1344 collect_moe.py"} - {framework: vllm, kernel_source: vllm_gptossmxfp4moe_marlin_marlinexperts, backend: marlin, source: "PR#1344 collect_moe.py"} diff --git a/python/aisimulate/collector/network/slurm/collect_trtllm_alltoall.py b/python/aisimulate/collector/network/slurm/collect_trtllm_alltoall.py index 1225fa97c7..1d6b868ab0 100644 --- a/python/aisimulate/collector/network/slurm/collect_trtllm_alltoall.py +++ b/python/aisimulate/collector/network/slurm/collect_trtllm_alltoall.py @@ -679,27 +679,80 @@ class AlltoallBenchmarkResult: combine_low_precision_error: Optional[BaseException] = field(default=None, compare=False) +def _agree_graph_capturable(func, label, device) -> None: + """Capture ``func`` once on every rank and agree on the outcome BEFORE any + rank replays a graph. + + benchmark_with_power captures locally and then replays the graph for its + warmups. These kernels are collectives: had capture failed on one rank + while its peers' capture succeeded, the peers' first replay would wait for + a transfer the failed rank never issues and the job would hang instead of + failing. The probe runs after one eager call (lazy kernel/workspace + initialisation is not capturable) and makes the capture outcome symmetric: + every rank raises, the run loop records the classified failure and all + ranks skip the case together — the agreement the loop already performs for + measured cases. The moe_a2a row has no column that could record an + eager-timed measurement, so the eager fallback used until 2026-10-12 is + gone (layer_permissions.md: measurement-method degradation is allowed only + when recorded in the output row). + """ + import torch + import torch.distributed as dist + + func() + torch.cuda.synchronize() + capture_error: BaseException | None = None + probe = torch.cuda.CUDAGraph() + try: + with torch.cuda.graph(probe): + func() + torch.cuda.synchronize() + except Exception as error: + capture_error = error + finally: + del probe + + failed_here = capture_error is not None + failed_anywhere = failed_here + if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1: + agreement = torch.tensor([1 if failed_here else 0], dtype=torch.int32, device=device) + dist.all_reduce(agreement, op=dist.ReduceOp.MAX) + failed_anywhere = bool(int(agreement.item())) + if failed_anywhere: + detail = f"{type(capture_error).__name__}: {capture_error}" if failed_here else "capture failed on a peer rank" + raise TrtllmAlltoallBenchmarkError( + f"CUDA graph capture of {label} failed ({detail}); the case fails classified instead of " + "shipping an eager launch-loop floor the moe_a2a row could not label" + ) + + def _benchmark_op(func, label, device, ep_rank, num_warmup, num_iterations): """Benchmark a single operation using shared benchmark_with_power. + CUDA-graph timed only: ``_agree_graph_capturable`` runs first so a capture + failure fails the case on every rank at once (see its docstring). + ``measure_power=False``: the moe_a2a table emits NO power column BY RULING (Task 4 review, see ``collect_moe_a2a._power_columns``): the ruling is per TABLE — HT/LL and trtllm writers share one file and one header, and ``helper.log_perf`` writes the header from the first row — so this writer must not sample or emit power either. """ + _agree_graph_capturable(func, label, device) with benchmark_with_power( device=device, kernel_func=func, num_warmups=num_warmup, num_runs=num_iterations, measure_power=False, - allow_graph_fail=True, ) as results: latency = results["latency_ms"] + if not results.get("used_cuda_graph", False): + raise TrtllmAlltoallBenchmarkError( + f"{label} was not timed under CUDA-graph replay; refusing to record an eager measurement" + ) if ep_rank == 0: - mode = "CUDA Graph" if results["used_cuda_graph"] else "Eager" - print(f" [{label}] {mode} timing: {latency:.3f} ms") + print(f" [{label}] CUDA Graph timing: {latency:.3f} ms") return latency diff --git a/python/aisimulate/collector/op_backend_catalog.yaml b/python/aisimulate/collector/op_backend_catalog.yaml index 055014b31b..ee0ccd3b86 100644 --- a/python/aisimulate/collector/op_backend_catalog.yaml +++ b/python/aisimulate/collector/op_backend_catalog.yaml @@ -118,12 +118,8 @@ families: op_files: - dsa_context_module_perf - dsa_generation_module_perf - - dsa_context_module_skip_indexer_perf - - dsa_generation_module_skip_indexer_perf - dsv4_csa_context_module_perf - dsv4_csa_generation_module_perf - - dsv4_csa_attn_module_perf - - dsv4_hca_attn_module_perf - dsv4_hca_context_module_perf - dsv4_hca_generation_module_perf - dsv4_paged_mqa_logits_module_perf diff --git a/python/aisimulate/collector/op_backend_facts.yaml b/python/aisimulate/collector/op_backend_facts.yaml index 0b15378cee..5d423cd26f 100644 --- a/python/aisimulate/collector/op_backend_facts.yaml +++ b/python/aisimulate/collector/op_backend_facts.yaml @@ -210,14 +210,14 @@ ops: - {framework: "vllm", version: "0.24.0", system: "l40s", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["flashinfer", "triton"], kernel_sources: ["vllm_flashinfer_fidecode", "vllm_flashinfer_fiprefill", "vllm_triton_attn"]} - {framework: "vllm", version: "0.24.0", system: "rtx_pro_6000_server", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2"], kernel_sources: ["vllm_flash_attn_fa2"]} - {framework: "vllm", version: "0.24.0", system: "rtx_pro_6000_server", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["flashinfer", "triton"], kernel_sources: ["vllm_flashinfer_fidecode", "vllm_flashinfer_fiprefill", "vllm_triton_attn"]} - - {framework: "vllm", version: "0.25.0", system: "b200_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["trtllm_mha", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_flashinfer_trtllmprefill"]} - - {framework: "vllm", version: "0.25.0", system: "b200_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_flashinfer_trtllmprefill", "vllm_triton_attn"]} - - {framework: "vllm", version: "0.25.0", system: "b300_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["trtllm_mha", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_flashinfer_trtllmprefill"]} - - {framework: "vllm", version: "0.25.0", system: "b300_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_flashinfer_trtllmprefill", "vllm_triton_attn"]} - - {framework: "vllm", version: "0.25.0", system: "gb200", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["trtllm_mha", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_flashinfer_trtllmprefill"]} - - {framework: "vllm", version: "0.25.0", system: "gb200", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_flashinfer_trtllmprefill", "vllm_triton_attn"]} - - {framework: "vllm", version: "0.25.0", system: "gb300", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["trtllm_mha", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_flashinfer_trtllmprefill"]} - - {framework: "vllm", version: "0.25.0", system: "gb300", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_flashinfer_trtllmprefill", "vllm_triton_attn"]} + - {framework: "vllm", version: "0.25.0", system: "b200_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_flashinfer_trtllmprefill"]} + - {framework: "vllm", version: "0.25.0", system: "b200_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_flashinfer_trtllmprefill", "vllm_triton_attn"]} + - {framework: "vllm", version: "0.25.0", system: "b300_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_flashinfer_trtllmprefill"]} + - {framework: "vllm", version: "0.25.0", system: "b300_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_flashinfer_trtllmprefill", "vllm_triton_attn"]} + - {framework: "vllm", version: "0.25.0", system: "gb200", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_flashinfer_trtllmprefill"]} + - {framework: "vllm", version: "0.25.0", system: "gb200", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_flashinfer_trtllmprefill", "vllm_triton_attn"]} + - {framework: "vllm", version: "0.25.0", system: "gb300", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_flashinfer_trtllmprefill"]} + - {framework: "vllm", version: "0.25.0", system: "gb300", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_flashinfer_trtllmprefill", "vllm_triton_attn"]} - {framework: "vllm", version: "0.25.0", system: "h100_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa3", "fa4"], kernel_sources: ["vllm_flash_attn_fa3", "vllm_flash_attn_fa4"]} - {framework: "vllm", version: "0.25.0", system: "h100_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["fa3"], kernel_sources: ["vllm_flash_attn_fa3"]} - {framework: "vllm", version: "0.25.0", system: "h200_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa3", "fa4"], kernel_sources: ["vllm_flash_attn_fa3", "vllm_flash_attn_fa4"]} @@ -226,14 +226,14 @@ ops: - {framework: "vllm", version: "0.26.0", system: "b60", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["flash_attn"], kernel_sources: ["vllm_flash_attn"]} - {framework: "vllm", version: "0.28.0", system: "b60", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["flash_attn"], kernel_sources: ["vllm_flash_attn"]} - {framework: "vllm", version: "0.28.0", system: "b60", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["flash_attn"], kernel_sources: ["vllm_flash_attn"]} - - {framework: "vllm", version: "0.30.0", system: "b200_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2", "trtllm_mha", "unverified"], kernel_sources: ["vllm_flash_attn_fa2", "vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_flashinfer_trtllmprefill"]} - - {framework: "vllm", version: "0.30.0", system: "b200_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_flashinfer_trtllmprefill", "vllm_triton_attn"]} - - {framework: "vllm", version: "0.30.0", system: "b300_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2", "trtllm_mha", "unverified"], kernel_sources: ["vllm_flash_attn_fa2", "vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_flashinfer_trtllmprefill"]} - - {framework: "vllm", version: "0.30.0", system: "b300_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_flashinfer_trtllmprefill", "vllm_triton_attn"]} - - {framework: "vllm", version: "0.30.0", system: "gb200", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2", "trtllm_mha", "unverified"], kernel_sources: ["vllm_flash_attn_fa2", "vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_flashinfer_trtllmprefill"]} - - {framework: "vllm", version: "0.30.0", system: "gb200", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_flashinfer_trtllmprefill", "vllm_triton_attn"]} - - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2", "trtllm_mha", "unverified"], kernel_sources: ["vllm_flash_attn_fa2", "vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_flashinfer_trtllmprefill"]} - - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_flashinfer_trtllmprefill", "vllm_triton_attn"]} + - {framework: "vllm", version: "0.30.0", system: "b200_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2", "trtllm_mha"], kernel_sources: ["vllm_flash_attn_fa2", "vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_flashinfer_trtllmprefill"]} + - {framework: "vllm", version: "0.30.0", system: "b200_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_flashinfer_trtllmprefill", "vllm_triton_attn"]} + - {framework: "vllm", version: "0.30.0", system: "b300_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2", "trtllm_mha"], kernel_sources: ["vllm_flash_attn_fa2", "vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_flashinfer_trtllmprefill"]} + - {framework: "vllm", version: "0.30.0", system: "b300_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_flashinfer_trtllmprefill", "vllm_triton_attn"]} + - {framework: "vllm", version: "0.30.0", system: "gb200", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2", "trtllm_mha"], kernel_sources: ["vllm_flash_attn_fa2", "vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_flashinfer_trtllmprefill"]} + - {framework: "vllm", version: "0.30.0", system: "gb200", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_flashinfer_trtllmprefill", "vllm_triton_attn"]} + - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2", "trtllm_mha"], kernel_sources: ["vllm_flash_attn_fa2", "vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_flashinfer_trtllmprefill"]} + - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_flashinfer_trtllmprefill", "vllm_triton_attn"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa3", "fa4"], kernel_sources: ["vllm_flash_attn_fa3", "vllm_flash_attn_fa4"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["fa3", "triton"], kernel_sources: ["vllm_flash_attn_fa3", "vllm_triton_attn"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa3", "fa4"], kernel_sources: ["vllm_flash_attn_fa3", "vllm_flash_attn_fa4"]} @@ -241,7 +241,7 @@ ops: - {framework: "vllm", version: "0.30.0", system: "l40s", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2", "flashinfer"], kernel_sources: ["vllm_flash_attn_fa2", "vllm_flashinfer_fidecode", "vllm_flashinfer_fiprefill"]} - {framework: "vllm", version: "0.30.0", system: "l40s", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["flashinfer", "triton"], kernel_sources: ["vllm_flashinfer_fidecode", "vllm_flashinfer_fiprefill", "vllm_triton_attn"]} - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2", "flashinfer"], kernel_sources: ["vllm_flash_attn_fa2", "vllm_flashinfer_fidecode", "vllm_flashinfer_fiprefill"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["flashinfer", "triton", "unverified"], kernel_sources: ["vllm_flashinfer_fiprefill", "vllm_flashinfer_flashinfertrtllmapidecode_xqa", "vllm_triton_attn"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "context_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["flashinfer", "triton", "trtllm_mha"], kernel_sources: ["vllm_flashinfer_fiprefill", "vllm_flashinfer_flashinfertrtllmapidecode_xqa", "vllm_triton_attn"]} - op_file: context_mla_perf family: mla axes: [op_name, kv_cache_dtype, mla_dtype] @@ -856,54 +856,54 @@ ops: - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "nvfp4", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "nvfp4", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "nvfp4", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "nvfp4", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "nvfp4", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "nvfp4", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_context_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "nvfp4", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - op_file: dsa_generation_module_perf family: sparse_attention axes: [op_name, architecture, kv_cache_dtype, mla_dtype, gemm_type] @@ -1366,54 +1366,54 @@ ops: - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "nvfp4", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} + - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flash_attn"], kernel_sources: ["FLASH_ATTN_MLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "nvfp4", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "nvfp4", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "nvfp4", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module", architecture: "DeepseekV32ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "nvfp4", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "nvfp4", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "bfloat16", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsa_generation_module_skip_indexer", architecture: "GlmMoeDsaForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "nvfp4", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_SM120"]} - op_file: dsv41_module_perf family: dsv41 axes: [] @@ -1478,7 +1478,7 @@ ops: - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "dsv4_csa_context_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsv4_csa_context_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsv4_csa_context_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsv4_csa_context_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_DSV4"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsv4_csa_context_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_DSV4"]} - op_file: dsv4_csa_generation_module_perf family: sparse_attention axes: [op_name, architecture, kv_cache_dtype, mla_dtype, gemm_type] @@ -1537,7 +1537,7 @@ ops: - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "dsv4_csa_generation_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsv4_csa_generation_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsv4_csa_generation_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsv4_csa_generation_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_DSV4"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsv4_csa_generation_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_DSV4"]} - op_file: dsv4_csa_topk_calib_perf family: sparse_attention axes: [op_name, architecture, kv_cache_dtype, mla_dtype, gemm_type] @@ -1554,22 +1554,6 @@ ops: - {framework: "sglang", version: "0.5.21", system: "gb300", op_name: "dsv4_csa_topk_calib", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8_e4m3", mla_dtype: "fp8_e4m3", gemm_type: "fp8_block", backends: ["sgl_kernel"], kernel_sources: ["topk_transform_v1", "topk_transform_v2"]} - {framework: "sglang", version: "0.5.21", system: "h100_sxm", op_name: "dsv4_csa_topk_calib", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8_e4m3", mla_dtype: "fp8_e4m3", gemm_type: "fp8_block", backends: ["sgl_kernel"], kernel_sources: ["topk_transform_v1", "topk_transform_v2"]} - {framework: "sglang", version: "0.5.21", system: "h200_sxm", op_name: "dsv4_csa_topk_calib", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8_e4m3", mla_dtype: "fp8_e4m3", gemm_type: "fp8_block", backends: ["sgl_kernel"], kernel_sources: ["topk_transform_v1", "topk_transform_v2"]} - - op_file: dsv4_hca_attn_module_perf - family: sparse_attention - axes: [op_name, architecture, kv_cache_dtype, mla_dtype, gemm_type] - facts: - - {framework: "vllm", version: "0.24.0", system: "b200_sxm", op_name: "dsv4_hca_attn_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - - {framework: "vllm", version: "0.24.0", system: "b300_sxm", op_name: "dsv4_hca_attn_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - - {framework: "vllm", version: "0.24.0", system: "gb200", op_name: "dsv4_hca_attn_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - - {framework: "vllm", version: "0.24.0", system: "gb300", op_name: "dsv4_hca_attn_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - - {framework: "vllm", version: "0.24.0", system: "h100_sxm", op_name: "dsv4_hca_attn_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - - {framework: "vllm", version: "0.24.0", system: "h200_sxm", op_name: "dsv4_hca_attn_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - - {framework: "vllm", version: "0.25.0", system: "b200_sxm", op_name: "dsv4_hca_attn_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - - {framework: "vllm", version: "0.25.0", system: "b300_sxm", op_name: "dsv4_hca_attn_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - - {framework: "vllm", version: "0.25.0", system: "gb200", op_name: "dsv4_hca_attn_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - - {framework: "vllm", version: "0.25.0", system: "gb300", op_name: "dsv4_hca_attn_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - - {framework: "vllm", version: "0.25.0", system: "h100_sxm", op_name: "dsv4_hca_attn_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - - {framework: "vllm", version: "0.25.0", system: "h200_sxm", op_name: "dsv4_hca_attn_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - op_file: dsv4_hca_context_module_perf family: sparse_attention axes: [op_name, architecture, kv_cache_dtype, mla_dtype, gemm_type] @@ -1630,7 +1614,7 @@ ops: - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "dsv4_hca_context_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsv4_hca_context_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsv4_hca_context_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsv4_hca_context_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_DSV4"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsv4_hca_context_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_DSV4"]} - op_file: dsv4_hca_generation_module_perf family: sparse_attention axes: [op_name, architecture, kv_cache_dtype, mla_dtype, gemm_type] @@ -1691,7 +1675,7 @@ ops: - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "dsv4_hca_generation_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "dsv4_hca_generation_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "dsv4_hca_generation_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashmla"], kernel_sources: ["FLASHMLA_SPARSE_DSV4"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsv4_hca_generation_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["unverified"], kernel_sources: ["FLASHINFER_MLA_SPARSE_DSV4"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "dsv4_hca_generation_module", architecture: "DeepseekV4ForCausalLM", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", gemm_type: "fp8_block", backends: ["flashinfer"], kernel_sources: ["FLASHINFER_MLA_SPARSE_DSV4"]} - op_file: dsv4_megamoe_module_perf family: moe axes: [op_name, phase, moe_dtype, kernel_dtype] @@ -1826,21 +1810,21 @@ ops: - {framework: "sglang", version: "0.5.17", system: "gb300", op_name: "gdn", phase: "context", backends: ["fla_triton", "unverified"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule"]} - {framework: "sglang", version: "0.5.17", system: "gb300", op_name: "gdn", phase: "generation", backends: ["fla_triton", "flashinfer", "unverified"], kernel_sources: ["causal_conv1d_update", "flashinfer_gated_delta_rule_decode", "fused_recurrent_gated_delta_rule_packed_decode"]} - {framework: "sglang", version: "0.5.21", system: "b200_sxm", op_name: "gdn", phase: "context", backends: ["fla_triton", "unverified"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule"]} - - {framework: "sglang", version: "0.5.21", system: "b200_sxm", op_name: "gdn", phase: "generation", backends: ["fla_triton", "unverified"], kernel_sources: ["fused_qkvzba_causal_conv1d_update_contiguous", "fused_recurrent_gated_delta_rule_packed_decode"]} + - {framework: "sglang", version: "0.5.21", system: "b200_sxm", op_name: "gdn", phase: "generation", backends: ["fla_triton", "triton"], kernel_sources: ["fused_qkvzba_causal_conv1d_update_contiguous", "fused_recurrent_gated_delta_rule_packed_decode"]} - {framework: "sglang", version: "0.5.21", system: "b300_sxm", op_name: "gdn", phase: "context", backends: ["fla_triton", "unverified"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule"]} - - {framework: "sglang", version: "0.5.21", system: "b300_sxm", op_name: "gdn", phase: "generation", backends: ["fla_triton", "unverified"], kernel_sources: ["fused_qkvzba_causal_conv1d_update_contiguous", "fused_recurrent_gated_delta_rule_packed_decode"]} + - {framework: "sglang", version: "0.5.21", system: "b300_sxm", op_name: "gdn", phase: "generation", backends: ["fla_triton", "triton"], kernel_sources: ["fused_qkvzba_causal_conv1d_update_contiguous", "fused_recurrent_gated_delta_rule_packed_decode"]} - {framework: "sglang", version: "0.5.21", system: "gb200", op_name: "gdn", phase: "context", backends: ["fla_triton", "unverified"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule"]} - - {framework: "sglang", version: "0.5.21", system: "gb200", op_name: "gdn", phase: "generation", backends: ["fla_triton", "unverified"], kernel_sources: ["fused_qkvzba_causal_conv1d_update_contiguous", "fused_recurrent_gated_delta_rule_packed_decode"]} + - {framework: "sglang", version: "0.5.21", system: "gb200", op_name: "gdn", phase: "generation", backends: ["fla_triton", "triton"], kernel_sources: ["fused_qkvzba_causal_conv1d_update_contiguous", "fused_recurrent_gated_delta_rule_packed_decode"]} - {framework: "sglang", version: "0.5.21", system: "gb300", op_name: "gdn", phase: "context", backends: ["fla_triton", "unverified"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule"]} - - {framework: "sglang", version: "0.5.21", system: "gb300", op_name: "gdn", phase: "generation", backends: ["fla_triton", "unverified"], kernel_sources: ["fused_qkvzba_causal_conv1d_update_contiguous", "fused_recurrent_gated_delta_rule_packed_decode"]} - - {framework: "sglang", version: "0.5.21", system: "h100_sxm", op_name: "gdn", phase: "context", backends: ["unverified"], kernel_sources: ["causal_conv1d_fn", "flashinfer_chunk_gated_delta_rule"]} - - {framework: "sglang", version: "0.5.21", system: "h100_sxm", op_name: "gdn", phase: "generation", backends: ["fla_triton", "unverified"], kernel_sources: ["fused_qkvzba_causal_conv1d_update_contiguous", "fused_recurrent_gated_delta_rule_packed_decode"]} - - {framework: "sglang", version: "0.5.21", system: "h200_sxm", op_name: "gdn", phase: "context", backends: ["unverified"], kernel_sources: ["causal_conv1d_fn", "flashinfer_chunk_gated_delta_rule"]} - - {framework: "sglang", version: "0.5.21", system: "h200_sxm", op_name: "gdn", phase: "generation", backends: ["fla_triton", "unverified"], kernel_sources: ["fused_qkvzba_causal_conv1d_update_contiguous", "fused_recurrent_gated_delta_rule_packed_decode"]} + - {framework: "sglang", version: "0.5.21", system: "gb300", op_name: "gdn", phase: "generation", backends: ["fla_triton", "triton"], kernel_sources: ["fused_qkvzba_causal_conv1d_update_contiguous", "fused_recurrent_gated_delta_rule_packed_decode"]} + - {framework: "sglang", version: "0.5.21", system: "h100_sxm", op_name: "gdn", phase: "context", backends: ["flashinfer", "unverified"], kernel_sources: ["causal_conv1d_fn", "flashinfer_chunk_gated_delta_rule"]} + - {framework: "sglang", version: "0.5.21", system: "h100_sxm", op_name: "gdn", phase: "generation", backends: ["fla_triton", "triton"], kernel_sources: ["fused_qkvzba_causal_conv1d_update_contiguous", "fused_recurrent_gated_delta_rule_packed_decode"]} + - {framework: "sglang", version: "0.5.21", system: "h200_sxm", op_name: "gdn", phase: "context", backends: ["flashinfer", "unverified"], kernel_sources: ["causal_conv1d_fn", "flashinfer_chunk_gated_delta_rule"]} + - {framework: "sglang", version: "0.5.21", system: "h200_sxm", op_name: "gdn", phase: "generation", backends: ["fla_triton", "triton"], kernel_sources: ["fused_qkvzba_causal_conv1d_update_contiguous", "fused_recurrent_gated_delta_rule_packed_decode"]} - {framework: "sglang", version: "0.5.21", system: "l40s", op_name: "gdn", phase: "context", backends: ["fla_triton", "unverified"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule"]} - - {framework: "sglang", version: "0.5.21", system: "l40s", op_name: "gdn", phase: "generation", backends: ["fla_triton", "unverified"], kernel_sources: ["fused_qkvzba_causal_conv1d_update_contiguous", "fused_recurrent_gated_delta_rule_packed_decode"]} + - {framework: "sglang", version: "0.5.21", system: "l40s", op_name: "gdn", phase: "generation", backends: ["fla_triton", "triton"], kernel_sources: ["fused_qkvzba_causal_conv1d_update_contiguous", "fused_recurrent_gated_delta_rule_packed_decode"]} - {framework: "sglang", version: "0.5.21", system: "rtx_pro_6000_server", op_name: "gdn", phase: "context", backends: ["fla_triton", "unverified"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule"]} - - {framework: "sglang", version: "0.5.21", system: "rtx_pro_6000_server", op_name: "gdn", phase: "generation", backends: ["fla_triton", "unverified"], kernel_sources: ["fused_qkvzba_causal_conv1d_update_contiguous", "fused_recurrent_gated_delta_rule_packed_decode"]} + - {framework: "sglang", version: "0.5.21", system: "rtx_pro_6000_server", op_name: "gdn", phase: "generation", backends: ["fla_triton", "triton"], kernel_sources: ["fused_qkvzba_causal_conv1d_update_contiguous", "fused_recurrent_gated_delta_rule_packed_decode"]} - {framework: "trtllm", version: "1.3.0rc20", system: "b200_sxm", op_name: "gdn", phase: "context", backends: ["causal_conv1d", "fla_triton"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule"]} - {framework: "trtllm", version: "1.3.0rc20", system: "b200_sxm", op_name: "gdn", phase: "generation", backends: ["causal_conv1d", "fla_triton"], kernel_sources: ["causal_conv1d_update", "fused_recurrent_gated_delta_rule"]} - {framework: "trtllm", version: "1.3.0rc20", system: "b300_sxm", op_name: "gdn", phase: "context", backends: ["causal_conv1d", "fla_triton"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule"]} @@ -1857,22 +1841,22 @@ ops: - {framework: "trtllm", version: "1.3.0rc20", system: "l40s", op_name: "gdn", phase: "generation", backends: ["causal_conv1d", "fla_triton"], kernel_sources: ["causal_conv1d_update", "fused_recurrent_gated_delta_rule"]} - {framework: "trtllm", version: "1.3.0rc20", system: "rtx_pro_6000_server", op_name: "gdn", phase: "context", backends: ["causal_conv1d", "fla_triton"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule"]} - {framework: "trtllm", version: "1.3.0rc20", system: "rtx_pro_6000_server", op_name: "gdn", phase: "generation", backends: ["causal_conv1d", "fla_triton"], kernel_sources: ["causal_conv1d_update", "fused_recurrent_gated_delta_rule"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "b200_sxm", op_name: "gdn", phase: "context", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule_flashinfer"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "b200_sxm", op_name: "gdn", phase: "generation", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_update", "fused_sigmoid_gating_delta_rule_update_flashinfer"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "b300_sxm", op_name: "gdn", phase: "context", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule_flashinfer"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "b300_sxm", op_name: "gdn", phase: "generation", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_update", "fused_sigmoid_gating_delta_rule_update_flashinfer"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "gb200", op_name: "gdn", phase: "context", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule_flashinfer"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "gb200", op_name: "gdn", phase: "generation", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_update", "fused_sigmoid_gating_delta_rule_update_flashinfer"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "gb300", op_name: "gdn", phase: "context", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule_flashinfer"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "gb300", op_name: "gdn", phase: "generation", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_update", "fused_sigmoid_gating_delta_rule_update_flashinfer"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "h100_sxm", op_name: "gdn", phase: "context", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule_flashinfer"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "h100_sxm", op_name: "gdn", phase: "generation", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_update", "fused_sigmoid_gating_delta_rule_update_flashinfer"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "h200_sxm", op_name: "gdn", phase: "context", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule_flashinfer"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "h200_sxm", op_name: "gdn", phase: "generation", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_update", "fused_sigmoid_gating_delta_rule_update_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "b200_sxm", op_name: "gdn", phase: "context", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "b200_sxm", op_name: "gdn", phase: "generation", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_update", "fused_sigmoid_gating_delta_rule_update_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "b300_sxm", op_name: "gdn", phase: "context", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "b300_sxm", op_name: "gdn", phase: "generation", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_update", "fused_sigmoid_gating_delta_rule_update_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "gb200", op_name: "gdn", phase: "context", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "gb200", op_name: "gdn", phase: "generation", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_update", "fused_sigmoid_gating_delta_rule_update_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "gb300", op_name: "gdn", phase: "context", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "gb300", op_name: "gdn", phase: "generation", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_update", "fused_sigmoid_gating_delta_rule_update_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "h100_sxm", op_name: "gdn", phase: "context", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "h100_sxm", op_name: "gdn", phase: "generation", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_update", "fused_sigmoid_gating_delta_rule_update_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "h200_sxm", op_name: "gdn", phase: "context", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "h200_sxm", op_name: "gdn", phase: "generation", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_update", "fused_sigmoid_gating_delta_rule_update_flashinfer"]} - {framework: "trtllm", version: "1.3.0rc29", system: "l40s", op_name: "gdn", phase: "context", backends: ["causal_conv1d", "fla_triton"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "l40s", op_name: "gdn", phase: "generation", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_update", "fused_sigmoid_gating_delta_rule_update_triton"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "rtx_pro_6000_server", op_name: "gdn", phase: "context", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule_flashinfer"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "rtx_pro_6000_server", op_name: "gdn", phase: "generation", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_update", "fused_sigmoid_gating_delta_rule_update_triton"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "l40s", op_name: "gdn", phase: "generation", backends: ["causal_conv1d", "fla_triton"], kernel_sources: ["causal_conv1d_update", "fused_sigmoid_gating_delta_rule_update_triton"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "rtx_pro_6000_server", op_name: "gdn", phase: "context", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "rtx_pro_6000_server", op_name: "gdn", phase: "generation", backends: ["causal_conv1d", "fla_triton"], kernel_sources: ["causal_conv1d_update", "fused_sigmoid_gating_delta_rule_update_triton"]} - {framework: "vllm", version: "0.24.0", system: "b200_sxm", op_name: "gdn", phase: "context", backends: ["flashinfer", "triton"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule_flashinfer"]} - {framework: "vllm", version: "0.24.0", system: "b200_sxm", op_name: "gdn", phase: "generation", backends: ["fla_triton", "triton"], kernel_sources: ["causal_conv1d_update", "fused_recurrent_gated_delta_rule_packed_decode"]} - {framework: "vllm", version: "0.24.0", system: "b300_sxm", op_name: "gdn", phase: "context", backends: ["flashinfer", "triton"], kernel_sources: ["causal_conv1d_fn", "chunk_gated_delta_rule_flashinfer"]} @@ -2176,7 +2160,7 @@ ops: - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "gemm", gemm_dtype: "fp8_block", backends: ["cutlass", "deepgemm", "flashinfer"], kernel_sources: ["CutlassFp8BlockScaledMMKernel", "DeepGemmFp8BlockScaledMMKernel", "FlashInferFp8BlockScaledMMKernel"]} - {framework: "vllm", version: "0.30.0", system: "l40s", op_name: "gemm", gemm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["torch.nn.functional.linear"]} - {framework: "vllm", version: "0.30.0", system: "l40s", op_name: "gemm", gemm_dtype: "fp8", backends: ["cutlass"], kernel_sources: ["CutlassFP8ScaledMMLinearKernel"]} - - {framework: "vllm", version: "0.30.0", system: "l40s", op_name: "gemm", gemm_dtype: "fp8_block", backends: ["unverified"], kernel_sources: ["MarlinFP8ScaledMMLinearKernel"]} + - {framework: "vllm", version: "0.30.0", system: "l40s", op_name: "gemm", gemm_dtype: "fp8_block", backends: ["marlin"], kernel_sources: ["MarlinFP8ScaledMMLinearKernel"]} - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "gemm", gemm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["torch.nn.functional.linear"]} - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "gemm", gemm_dtype: "fp8", backends: ["cutlass"], kernel_sources: ["CutlassFP8ScaledMMLinearKernel"]} - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "gemm", gemm_dtype: "fp8_block", backends: ["cutlass", "deepgemm"], kernel_sources: ["CutlassFp8BlockScaledMMKernel", "DeepGemmFp8BlockScaledMMKernel"]} @@ -2288,14 +2272,14 @@ ops: - {framework: "vllm", version: "0.24.0", system: "l40s", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["flashinfer", "triton"], kernel_sources: ["vllm_flashinfer_fidecode", "vllm_triton_attn"]} - {framework: "vllm", version: "0.24.0", system: "rtx_pro_6000_server", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2"], kernel_sources: ["vllm_flash_attn_fa2"]} - {framework: "vllm", version: "0.24.0", system: "rtx_pro_6000_server", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["flashinfer", "triton"], kernel_sources: ["vllm_flashinfer_fidecode", "vllm_triton_attn"]} - - {framework: "vllm", version: "0.25.0", system: "b200_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode"]} - - {framework: "vllm", version: "0.25.0", system: "b200_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_triton_attn"]} - - {framework: "vllm", version: "0.25.0", system: "b300_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode"]} - - {framework: "vllm", version: "0.25.0", system: "b300_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_triton_attn"]} - - {framework: "vllm", version: "0.25.0", system: "gb200", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode"]} - - {framework: "vllm", version: "0.25.0", system: "gb200", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_triton_attn"]} - - {framework: "vllm", version: "0.25.0", system: "gb300", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode"]} - - {framework: "vllm", version: "0.25.0", system: "gb300", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_triton_attn"]} + - {framework: "vllm", version: "0.25.0", system: "b200_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode"]} + - {framework: "vllm", version: "0.25.0", system: "b200_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_triton_attn"]} + - {framework: "vllm", version: "0.25.0", system: "b300_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode"]} + - {framework: "vllm", version: "0.25.0", system: "b300_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_triton_attn"]} + - {framework: "vllm", version: "0.25.0", system: "gb200", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode"]} + - {framework: "vllm", version: "0.25.0", system: "gb200", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_triton_attn"]} + - {framework: "vllm", version: "0.25.0", system: "gb300", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode"]} + - {framework: "vllm", version: "0.25.0", system: "gb300", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode", "vllm_triton_attn"]} - {framework: "vllm", version: "0.25.0", system: "h100_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa3", "fa4"], kernel_sources: ["vllm_flash_attn_fa3", "vllm_flash_attn_fa4"]} - {framework: "vllm", version: "0.25.0", system: "h100_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["fa3"], kernel_sources: ["vllm_flash_attn_fa3"]} - {framework: "vllm", version: "0.25.0", system: "h200_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa3", "fa4"], kernel_sources: ["vllm_flash_attn_fa3", "vllm_flash_attn_fa4"]} @@ -2304,14 +2288,14 @@ ops: - {framework: "vllm", version: "0.26.0", system: "b60", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["flash_attn"], kernel_sources: ["vllm_flash_attn"]} - {framework: "vllm", version: "0.28.0", system: "b60", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["flash_attn"], kernel_sources: ["vllm_flash_attn"]} - {framework: "vllm", version: "0.28.0", system: "b60", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["flash_attn"], kernel_sources: ["vllm_flash_attn"]} - - {framework: "vllm", version: "0.30.0", system: "b200_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2", "unverified"], kernel_sources: ["vllm_flash_attn_fa2", "vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen"]} - - {framework: "vllm", version: "0.30.0", system: "b200_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_triton_attn"]} - - {framework: "vllm", version: "0.30.0", system: "b300_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2", "unverified"], kernel_sources: ["vllm_flash_attn_fa2", "vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen"]} - - {framework: "vllm", version: "0.30.0", system: "b300_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_triton_attn"]} - - {framework: "vllm", version: "0.30.0", system: "gb200", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2", "unverified"], kernel_sources: ["vllm_flash_attn_fa2", "vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen"]} - - {framework: "vllm", version: "0.30.0", system: "gb200", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_triton_attn"]} - - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2", "unverified"], kernel_sources: ["vllm_flash_attn_fa2", "vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen"]} - - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_triton_attn"]} + - {framework: "vllm", version: "0.30.0", system: "b200_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2", "trtllm_mha"], kernel_sources: ["vllm_flash_attn_fa2", "vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen"]} + - {framework: "vllm", version: "0.30.0", system: "b200_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_triton_attn"]} + - {framework: "vllm", version: "0.30.0", system: "b300_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2", "trtllm_mha"], kernel_sources: ["vllm_flash_attn_fa2", "vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen"]} + - {framework: "vllm", version: "0.30.0", system: "b300_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_triton_attn"]} + - {framework: "vllm", version: "0.30.0", system: "gb200", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2", "trtllm_mha"], kernel_sources: ["vllm_flash_attn_fa2", "vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen"]} + - {framework: "vllm", version: "0.30.0", system: "gb200", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_triton_attn"]} + - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2", "trtllm_mha"], kernel_sources: ["vllm_flash_attn_fa2", "vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen"]} + - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode_trtllm_gen", "vllm_triton_attn"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa3", "fa4"], kernel_sources: ["vllm_flash_attn_fa3", "vllm_flash_attn_fa4"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["fa3", "triton"], kernel_sources: ["vllm_flash_attn_fa3", "vllm_triton_attn"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa3", "fa4"], kernel_sources: ["vllm_flash_attn_fa3", "vllm_flash_attn_fa4"]} @@ -2319,7 +2303,7 @@ ops: - {framework: "vllm", version: "0.30.0", system: "l40s", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2", "flashinfer"], kernel_sources: ["vllm_flash_attn_fa2", "vllm_flashinfer_fidecode"]} - {framework: "vllm", version: "0.30.0", system: "l40s", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["flashinfer", "triton"], kernel_sources: ["vllm_flashinfer_fidecode", "vllm_triton_attn"]} - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "bfloat16", backends: ["fa2", "flashinfer"], kernel_sources: ["vllm_flash_attn_fa2", "vllm_flashinfer_fidecode"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "unverified"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode_xqa", "vllm_triton_attn"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "generation_attention", attn_dtype: "bfloat16", kv_cache_dtype: "fp8", backends: ["triton", "trtllm_mha"], kernel_sources: ["vllm_flashinfer_flashinfertrtllmapidecode_xqa", "vllm_triton_attn"]} - op_file: generation_mla_perf family: mla axes: [op_name, kv_cache_dtype, mla_dtype] @@ -2541,22 +2525,22 @@ ops: - {framework: "trtllm", version: "1.3.0rc20", system: "l40s", op_name: "mamba2", phase: "generation", backends: ["causal_conv1d"], kernel_sources: ["causal_conv1d_update"]} - {framework: "trtllm", version: "1.3.0rc20", system: "rtx_pro_6000_server", op_name: "mamba2", phase: "context", backends: ["causal_conv1d"], kernel_sources: ["causal_conv1d_fn"]} - {framework: "trtllm", version: "1.3.0rc20", system: "rtx_pro_6000_server", op_name: "mamba2", phase: "generation", backends: ["causal_conv1d"], kernel_sources: ["causal_conv1d_update"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "b200_sxm", op_name: "mamba2", phase: "context", backends: ["flashinfer", "unverified"], kernel_sources: ["causal_conv1d_fn_token_major", "mamba_chunk_scan_combined"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "b200_sxm", op_name: "mamba2", phase: "generation", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_update", "selective_state_update_flashinfer"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "b300_sxm", op_name: "mamba2", phase: "context", backends: ["flashinfer", "unverified"], kernel_sources: ["causal_conv1d_fn_token_major", "mamba_chunk_scan_combined"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "b300_sxm", op_name: "mamba2", phase: "generation", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_update", "selective_state_update_flashinfer"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "gb200", op_name: "mamba2", phase: "context", backends: ["flashinfer", "unverified"], kernel_sources: ["causal_conv1d_fn_token_major", "mamba_chunk_scan_combined"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "gb200", op_name: "mamba2", phase: "generation", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_update", "selective_state_update_flashinfer"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "gb300", op_name: "mamba2", phase: "context", backends: ["flashinfer", "unverified"], kernel_sources: ["causal_conv1d_fn_token_major", "mamba_chunk_scan_combined"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "gb300", op_name: "mamba2", phase: "generation", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_update", "selective_state_update_flashinfer"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "h100_sxm", op_name: "mamba2", phase: "context", backends: ["flashinfer", "unverified"], kernel_sources: ["causal_conv1d_fn_token_major", "mamba_chunk_scan_combined"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "h100_sxm", op_name: "mamba2", phase: "generation", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_update", "selective_state_update_flashinfer"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "h200_sxm", op_name: "mamba2", phase: "context", backends: ["flashinfer", "unverified"], kernel_sources: ["causal_conv1d_fn_token_major", "mamba_chunk_scan_combined"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "h200_sxm", op_name: "mamba2", phase: "generation", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_update", "selective_state_update_flashinfer"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "l40s", op_name: "mamba2", phase: "context", backends: ["flashinfer", "unverified"], kernel_sources: ["causal_conv1d_fn_token_major", "mamba_chunk_scan_combined"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "l40s", op_name: "mamba2", phase: "generation", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_update", "selective_state_update_flashinfer"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "rtx_pro_6000_server", op_name: "mamba2", phase: "context", backends: ["flashinfer", "unverified"], kernel_sources: ["causal_conv1d_fn_token_major", "mamba_chunk_scan_combined"]} - - {framework: "trtllm", version: "1.3.0rc29", system: "rtx_pro_6000_server", op_name: "mamba2", phase: "generation", backends: ["causal_conv1d", "unverified"], kernel_sources: ["causal_conv1d_update", "selective_state_update_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "b200_sxm", op_name: "mamba2", phase: "context", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_fn_token_major", "mamba_chunk_scan_combined"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "b200_sxm", op_name: "mamba2", phase: "generation", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_update", "selective_state_update_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "b300_sxm", op_name: "mamba2", phase: "context", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_fn_token_major", "mamba_chunk_scan_combined"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "b300_sxm", op_name: "mamba2", phase: "generation", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_update", "selective_state_update_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "gb200", op_name: "mamba2", phase: "context", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_fn_token_major", "mamba_chunk_scan_combined"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "gb200", op_name: "mamba2", phase: "generation", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_update", "selective_state_update_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "gb300", op_name: "mamba2", phase: "context", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_fn_token_major", "mamba_chunk_scan_combined"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "gb300", op_name: "mamba2", phase: "generation", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_update", "selective_state_update_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "h100_sxm", op_name: "mamba2", phase: "context", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_fn_token_major", "mamba_chunk_scan_combined"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "h100_sxm", op_name: "mamba2", phase: "generation", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_update", "selective_state_update_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "h200_sxm", op_name: "mamba2", phase: "context", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_fn_token_major", "mamba_chunk_scan_combined"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "h200_sxm", op_name: "mamba2", phase: "generation", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_update", "selective_state_update_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "l40s", op_name: "mamba2", phase: "context", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_fn_token_major", "mamba_chunk_scan_combined"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "l40s", op_name: "mamba2", phase: "generation", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_update", "selective_state_update_flashinfer"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "rtx_pro_6000_server", op_name: "mamba2", phase: "context", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_fn_token_major", "mamba_chunk_scan_combined"]} + - {framework: "trtllm", version: "1.3.0rc29", system: "rtx_pro_6000_server", op_name: "mamba2", phase: "generation", backends: ["causal_conv1d", "flashinfer"], kernel_sources: ["causal_conv1d_update", "selective_state_update_flashinfer"]} - op_file: mhc_module_perf family: mhc axes: [op_name, architecture] @@ -2746,37 +2730,37 @@ ops: - {framework: "sglang", version: "0.5.14", system: "rtx_pro_6000_server", op_name: "mla_gen_pre", bmm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["sglang_torch_bmm"]} - {framework: "sglang", version: "0.5.14", system: "rtx_pro_6000_server", op_name: "mla_gen_pre", bmm_dtype: "fp8", backends: ["cublas"], kernel_sources: ["sglang_sgl_kernel_bmm_fp8"]} - {framework: "sglang", version: "0.5.21", system: "b200_sxm", op_name: "mla_gen_post", bmm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["sglang_torch_bmm"]} - - {framework: "sglang", version: "0.5.21", system: "b200_sxm", op_name: "mla_gen_post", bmm_dtype: "fp8", backends: ["unverified"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} + - {framework: "sglang", version: "0.5.21", system: "b200_sxm", op_name: "mla_gen_post", bmm_dtype: "fp8", backends: ["cublas"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} - {framework: "sglang", version: "0.5.21", system: "b200_sxm", op_name: "mla_gen_pre", bmm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["sglang_torch_bmm"]} - - {framework: "sglang", version: "0.5.21", system: "b200_sxm", op_name: "mla_gen_pre", bmm_dtype: "fp8", backends: ["unverified"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} + - {framework: "sglang", version: "0.5.21", system: "b200_sxm", op_name: "mla_gen_pre", bmm_dtype: "fp8", backends: ["cublas"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} - {framework: "sglang", version: "0.5.21", system: "b300_sxm", op_name: "mla_gen_post", bmm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["sglang_torch_bmm"]} - - {framework: "sglang", version: "0.5.21", system: "b300_sxm", op_name: "mla_gen_post", bmm_dtype: "fp8", backends: ["unverified"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} + - {framework: "sglang", version: "0.5.21", system: "b300_sxm", op_name: "mla_gen_post", bmm_dtype: "fp8", backends: ["cublas"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} - {framework: "sglang", version: "0.5.21", system: "b300_sxm", op_name: "mla_gen_pre", bmm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["sglang_torch_bmm"]} - - {framework: "sglang", version: "0.5.21", system: "b300_sxm", op_name: "mla_gen_pre", bmm_dtype: "fp8", backends: ["unverified"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} + - {framework: "sglang", version: "0.5.21", system: "b300_sxm", op_name: "mla_gen_pre", bmm_dtype: "fp8", backends: ["cublas"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} - {framework: "sglang", version: "0.5.21", system: "gb200", op_name: "mla_gen_post", bmm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["sglang_torch_bmm"]} - - {framework: "sglang", version: "0.5.21", system: "gb200", op_name: "mla_gen_post", bmm_dtype: "fp8", backends: ["unverified"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} + - {framework: "sglang", version: "0.5.21", system: "gb200", op_name: "mla_gen_post", bmm_dtype: "fp8", backends: ["cublas"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} - {framework: "sglang", version: "0.5.21", system: "gb200", op_name: "mla_gen_pre", bmm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["sglang_torch_bmm"]} - - {framework: "sglang", version: "0.5.21", system: "gb200", op_name: "mla_gen_pre", bmm_dtype: "fp8", backends: ["unverified"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} + - {framework: "sglang", version: "0.5.21", system: "gb200", op_name: "mla_gen_pre", bmm_dtype: "fp8", backends: ["cublas"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} - {framework: "sglang", version: "0.5.21", system: "gb300", op_name: "mla_gen_post", bmm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["sglang_torch_bmm"]} - - {framework: "sglang", version: "0.5.21", system: "gb300", op_name: "mla_gen_post", bmm_dtype: "fp8", backends: ["unverified"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} + - {framework: "sglang", version: "0.5.21", system: "gb300", op_name: "mla_gen_post", bmm_dtype: "fp8", backends: ["cublas"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} - {framework: "sglang", version: "0.5.21", system: "gb300", op_name: "mla_gen_pre", bmm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["sglang_torch_bmm"]} - - {framework: "sglang", version: "0.5.21", system: "gb300", op_name: "mla_gen_pre", bmm_dtype: "fp8", backends: ["unverified"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} + - {framework: "sglang", version: "0.5.21", system: "gb300", op_name: "mla_gen_pre", bmm_dtype: "fp8", backends: ["cublas"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} - {framework: "sglang", version: "0.5.21", system: "h100_sxm", op_name: "mla_gen_post", bmm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["sglang_torch_bmm"]} - - {framework: "sglang", version: "0.5.21", system: "h100_sxm", op_name: "mla_gen_post", bmm_dtype: "fp8", backends: ["unverified"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} + - {framework: "sglang", version: "0.5.21", system: "h100_sxm", op_name: "mla_gen_post", bmm_dtype: "fp8", backends: ["cublas"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} - {framework: "sglang", version: "0.5.21", system: "h100_sxm", op_name: "mla_gen_pre", bmm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["sglang_torch_bmm"]} - - {framework: "sglang", version: "0.5.21", system: "h100_sxm", op_name: "mla_gen_pre", bmm_dtype: "fp8", backends: ["unverified"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} + - {framework: "sglang", version: "0.5.21", system: "h100_sxm", op_name: "mla_gen_pre", bmm_dtype: "fp8", backends: ["cublas"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} - {framework: "sglang", version: "0.5.21", system: "h200_sxm", op_name: "mla_gen_post", bmm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["sglang_torch_bmm"]} - - {framework: "sglang", version: "0.5.21", system: "h200_sxm", op_name: "mla_gen_post", bmm_dtype: "fp8", backends: ["unverified"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} + - {framework: "sglang", version: "0.5.21", system: "h200_sxm", op_name: "mla_gen_post", bmm_dtype: "fp8", backends: ["cublas"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} - {framework: "sglang", version: "0.5.21", system: "h200_sxm", op_name: "mla_gen_pre", bmm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["sglang_torch_bmm"]} - - {framework: "sglang", version: "0.5.21", system: "h200_sxm", op_name: "mla_gen_pre", bmm_dtype: "fp8", backends: ["unverified"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} + - {framework: "sglang", version: "0.5.21", system: "h200_sxm", op_name: "mla_gen_pre", bmm_dtype: "fp8", backends: ["cublas"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} - {framework: "sglang", version: "0.5.21", system: "l40s", op_name: "mla_gen_post", bmm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["sglang_torch_bmm"]} - - {framework: "sglang", version: "0.5.21", system: "l40s", op_name: "mla_gen_post", bmm_dtype: "fp8", backends: ["unverified"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} + - {framework: "sglang", version: "0.5.21", system: "l40s", op_name: "mla_gen_post", bmm_dtype: "fp8", backends: ["cublas"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} - {framework: "sglang", version: "0.5.21", system: "l40s", op_name: "mla_gen_pre", bmm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["sglang_torch_bmm"]} - - {framework: "sglang", version: "0.5.21", system: "l40s", op_name: "mla_gen_pre", bmm_dtype: "fp8", backends: ["unverified"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} + - {framework: "sglang", version: "0.5.21", system: "l40s", op_name: "mla_gen_pre", bmm_dtype: "fp8", backends: ["cublas"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} - {framework: "sglang", version: "0.5.21", system: "rtx_pro_6000_server", op_name: "mla_gen_post", bmm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["sglang_torch_bmm"]} - - {framework: "sglang", version: "0.5.21", system: "rtx_pro_6000_server", op_name: "mla_gen_post", bmm_dtype: "fp8", backends: ["unverified"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} + - {framework: "sglang", version: "0.5.21", system: "rtx_pro_6000_server", op_name: "mla_gen_post", bmm_dtype: "fp8", backends: ["cublas"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} - {framework: "sglang", version: "0.5.21", system: "rtx_pro_6000_server", op_name: "mla_gen_pre", bmm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["sglang_torch_bmm"]} - - {framework: "sglang", version: "0.5.21", system: "rtx_pro_6000_server", op_name: "mla_gen_pre", bmm_dtype: "fp8", backends: ["unverified"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} + - {framework: "sglang", version: "0.5.21", system: "rtx_pro_6000_server", op_name: "mla_gen_pre", bmm_dtype: "fp8", backends: ["cublas"], kernel_sources: ["sglang_flashinfer_bmm_fp8"]} - {framework: "trtllm", version: "1.0.0", system: "a100_sxm", op_name: "mla_gen_post", bmm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["default"]} - {framework: "trtllm", version: "1.0.0", system: "a100_sxm", op_name: "mla_gen_pre", bmm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["default"]} - {framework: "trtllm", version: "1.0.0", system: "l40s", op_name: "mla_gen_post", bmm_dtype: "bfloat16", backends: ["cublas"], kernel_sources: ["default"]} @@ -3693,50 +3677,50 @@ ops: - {framework: "vllm", version: "0.30.0", system: "b200_sxm", op_name: "moe", moe_dtype: "bfloat16", backends: ["flashinfer_cutlass", "flashinfer_trtllm_gen", "triton_fused_moe"], kernel_sources: ["vllm_unquantizedfusedmoe_flashinfer_cutlass_flashinferexperts", "vllm_unquantizedfusedmoe_flashinfer_trtllm_trtllmbf16expertsmonolithic", "vllm_unquantizedfusedmoe_triton_tritonexperts"]} - {framework: "vllm", version: "0.30.0", system: "b200_sxm", op_name: "moe", moe_dtype: "fp8", backends: ["flashinfer_cutlass", "flashinfer_trtllm_gen"], kernel_sources: ["vllm_modeloptfp8moe_flashinfer_cutlass_flashinferexperts", "vllm_modeloptfp8moe_flashinfer_trtllm_trtllmfp8expertsmonolithic"]} - {framework: "vllm", version: "0.30.0", system: "b200_sxm", op_name: "moe", moe_dtype: "fp8_block", backends: ["flashinfer_trtllm_gen", "triton_fused_moe"], kernel_sources: ["vllm_fp8moe_flashinfer_trtllm_trtllmfp8expertsmodular", "vllm_fp8moe_flashinfer_trtllm_trtllmfp8expertsmonolithic", "vllm_fp8moe_triton_tritonexperts"]} - - {framework: "vllm", version: "0.30.0", system: "b200_sxm", op_name: "moe", moe_dtype: "int4_wo", backends: ["unverified"], kernel_sources: ["vllm_compressedtensorswna16moe_marlin_marlinexperts"]} - - {framework: "vllm", version: "0.30.0", system: "b200_sxm", op_name: "moe", moe_dtype: "nvfp4", backends: ["flashinfer_trtllm_gen", "unverified"], kernel_sources: ["vllm_compressedtensorsw4a4nvfp4moe_flashinfer_cutedsl_flashinfercutedslexperts", "vllm_compressedtensorsw4a4nvfp4moe_flashinfer_trtllm_trtllmnvfp4expertsmodular", "vllm_compressedtensorsw4a4nvfp4moe_flashinfer_trtllm_trtllmnvfp4expertsmonolithic"]} + - {framework: "vllm", version: "0.30.0", system: "b200_sxm", op_name: "moe", moe_dtype: "int4_wo", backends: ["marlin"], kernel_sources: ["vllm_compressedtensorswna16moe_marlin_marlinexperts"]} + - {framework: "vllm", version: "0.30.0", system: "b200_sxm", op_name: "moe", moe_dtype: "nvfp4", backends: ["flashinfer_cutedsl", "flashinfer_trtllm_gen"], kernel_sources: ["vllm_compressedtensorsw4a4nvfp4moe_flashinfer_cutedsl_flashinfercutedslexperts", "vllm_compressedtensorsw4a4nvfp4moe_flashinfer_trtllm_trtllmnvfp4expertsmodular", "vllm_compressedtensorsw4a4nvfp4moe_flashinfer_trtllm_trtllmnvfp4expertsmonolithic"]} - {framework: "vllm", version: "0.30.0", system: "b200_sxm", op_name: "moe", moe_dtype: "w4a16_mxfp4", backends: ["flashinfer_trtllm_gen", "marlin"], kernel_sources: ["vllm_compressedtensorsw4a4mxfp4moe_marlin_marlinexperts_situ_as_silu", "vllm_gptossmxfp4moe_flashinfer_trtllm_mxfp4_bf16_trtllmmxfp4expertsmonolithic"]} - {framework: "vllm", version: "0.30.0", system: "b200_sxm", op_name: "moe", moe_dtype: "w4a8_mxfp4_mxfp8", backends: ["flashinfer_trtllm_gen"], kernel_sources: ["vllm_mxfp4moe_flashinfer_trtllm_mxfp4_mxfp8_trtllmmxfp4expertsmodular"]} - {framework: "vllm", version: "0.30.0", system: "b300_sxm", op_name: "moe", moe_dtype: "bfloat16", backends: ["flashinfer_cutlass", "flashinfer_trtllm_gen", "triton_fused_moe"], kernel_sources: ["vllm_unquantizedfusedmoe_flashinfer_cutlass_flashinferexperts", "vllm_unquantizedfusedmoe_flashinfer_trtllm_trtllmbf16expertsmonolithic", "vllm_unquantizedfusedmoe_triton_tritonexperts"]} - {framework: "vllm", version: "0.30.0", system: "b300_sxm", op_name: "moe", moe_dtype: "fp8", backends: ["flashinfer_cutlass", "flashinfer_trtllm_gen"], kernel_sources: ["vllm_modeloptfp8moe_flashinfer_cutlass_flashinferexperts", "vllm_modeloptfp8moe_flashinfer_trtllm_trtllmfp8expertsmonolithic"]} - {framework: "vllm", version: "0.30.0", system: "b300_sxm", op_name: "moe", moe_dtype: "fp8_block", backends: ["flashinfer_trtllm_gen", "triton_fused_moe"], kernel_sources: ["vllm_fp8moe_flashinfer_trtllm_trtllmfp8expertsmodular", "vllm_fp8moe_flashinfer_trtllm_trtllmfp8expertsmonolithic", "vllm_fp8moe_triton_tritonexperts"]} - - {framework: "vllm", version: "0.30.0", system: "b300_sxm", op_name: "moe", moe_dtype: "int4_wo", backends: ["unverified"], kernel_sources: ["vllm_compressedtensorswna16moe_marlin_marlinexperts"]} - - {framework: "vllm", version: "0.30.0", system: "b300_sxm", op_name: "moe", moe_dtype: "nvfp4", backends: ["flashinfer_trtllm_gen", "unverified"], kernel_sources: ["vllm_compressedtensorsw4a4nvfp4moe_flashinfer_cutedsl_flashinfercutedslexperts", "vllm_compressedtensorsw4a4nvfp4moe_flashinfer_trtllm_trtllmnvfp4expertsmodular", "vllm_compressedtensorsw4a4nvfp4moe_flashinfer_trtllm_trtllmnvfp4expertsmonolithic"]} + - {framework: "vllm", version: "0.30.0", system: "b300_sxm", op_name: "moe", moe_dtype: "int4_wo", backends: ["marlin"], kernel_sources: ["vllm_compressedtensorswna16moe_marlin_marlinexperts"]} + - {framework: "vllm", version: "0.30.0", system: "b300_sxm", op_name: "moe", moe_dtype: "nvfp4", backends: ["flashinfer_cutedsl", "flashinfer_trtllm_gen"], kernel_sources: ["vllm_compressedtensorsw4a4nvfp4moe_flashinfer_cutedsl_flashinfercutedslexperts", "vllm_compressedtensorsw4a4nvfp4moe_flashinfer_trtllm_trtllmnvfp4expertsmodular", "vllm_compressedtensorsw4a4nvfp4moe_flashinfer_trtllm_trtllmnvfp4expertsmonolithic"]} - {framework: "vllm", version: "0.30.0", system: "b300_sxm", op_name: "moe", moe_dtype: "w4a16_mxfp4", backends: ["flashinfer_trtllm_gen", "marlin"], kernel_sources: ["vllm_compressedtensorsw4a4mxfp4moe_marlin_marlinexperts_situ_as_silu", "vllm_gptossmxfp4moe_flashinfer_trtllm_mxfp4_bf16_trtllmmxfp4expertsmonolithic"]} - {framework: "vllm", version: "0.30.0", system: "b300_sxm", op_name: "moe", moe_dtype: "w4a8_mxfp4_mxfp8", backends: ["flashinfer_trtllm_gen"], kernel_sources: ["vllm_mxfp4moe_flashinfer_trtllm_mxfp4_mxfp8_trtllmmxfp4expertsmodular"]} - {framework: "vllm", version: "0.30.0", system: "gb200", op_name: "moe", moe_dtype: "bfloat16", backends: ["flashinfer_cutlass", "flashinfer_trtllm_gen", "triton_fused_moe"], kernel_sources: ["vllm_unquantizedfusedmoe_flashinfer_cutlass_flashinferexperts", "vllm_unquantizedfusedmoe_flashinfer_trtllm_trtllmbf16expertsmonolithic", "vllm_unquantizedfusedmoe_triton_tritonexperts"]} - {framework: "vllm", version: "0.30.0", system: "gb200", op_name: "moe", moe_dtype: "fp8", backends: ["flashinfer_cutlass", "flashinfer_trtllm_gen"], kernel_sources: ["vllm_modeloptfp8moe_flashinfer_cutlass_flashinferexperts", "vllm_modeloptfp8moe_flashinfer_trtllm_trtllmfp8expertsmonolithic"]} - {framework: "vllm", version: "0.30.0", system: "gb200", op_name: "moe", moe_dtype: "fp8_block", backends: ["flashinfer_trtllm_gen", "triton_fused_moe"], kernel_sources: ["vllm_fp8moe_flashinfer_trtllm_trtllmfp8expertsmodular", "vllm_fp8moe_flashinfer_trtllm_trtllmfp8expertsmonolithic", "vllm_fp8moe_triton_tritonexperts"]} - - {framework: "vllm", version: "0.30.0", system: "gb200", op_name: "moe", moe_dtype: "int4_wo", backends: ["unverified"], kernel_sources: ["vllm_compressedtensorswna16moe_marlin_marlinexperts"]} - - {framework: "vllm", version: "0.30.0", system: "gb200", op_name: "moe", moe_dtype: "nvfp4", backends: ["flashinfer_trtllm_gen", "unverified"], kernel_sources: ["vllm_compressedtensorsw4a4nvfp4moe_flashinfer_cutedsl_flashinfercutedslexperts", "vllm_compressedtensorsw4a4nvfp4moe_flashinfer_trtllm_trtllmnvfp4expertsmodular", "vllm_compressedtensorsw4a4nvfp4moe_flashinfer_trtllm_trtllmnvfp4expertsmonolithic"]} + - {framework: "vllm", version: "0.30.0", system: "gb200", op_name: "moe", moe_dtype: "int4_wo", backends: ["marlin"], kernel_sources: ["vllm_compressedtensorswna16moe_marlin_marlinexperts"]} + - {framework: "vllm", version: "0.30.0", system: "gb200", op_name: "moe", moe_dtype: "nvfp4", backends: ["flashinfer_cutedsl", "flashinfer_trtllm_gen"], kernel_sources: ["vllm_compressedtensorsw4a4nvfp4moe_flashinfer_cutedsl_flashinfercutedslexperts", "vllm_compressedtensorsw4a4nvfp4moe_flashinfer_trtllm_trtllmnvfp4expertsmodular", "vllm_compressedtensorsw4a4nvfp4moe_flashinfer_trtllm_trtllmnvfp4expertsmonolithic"]} - {framework: "vllm", version: "0.30.0", system: "gb200", op_name: "moe", moe_dtype: "w4a16_mxfp4", backends: ["flashinfer_trtllm_gen", "marlin"], kernel_sources: ["vllm_compressedtensorsw4a4mxfp4moe_marlin_marlinexperts_situ_as_silu", "vllm_gptossmxfp4moe_flashinfer_trtllm_mxfp4_bf16_trtllmmxfp4expertsmonolithic"]} - {framework: "vllm", version: "0.30.0", system: "gb200", op_name: "moe", moe_dtype: "w4a8_mxfp4_mxfp8", backends: ["flashinfer_trtllm_gen"], kernel_sources: ["vllm_mxfp4moe_flashinfer_trtllm_mxfp4_mxfp8_trtllmmxfp4expertsmodular"]} - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "moe", moe_dtype: "bfloat16", backends: ["flashinfer_cutlass", "flashinfer_trtllm_gen", "triton_fused_moe"], kernel_sources: ["vllm_unquantizedfusedmoe_flashinfer_cutlass_flashinferexperts", "vllm_unquantizedfusedmoe_flashinfer_trtllm_trtllmbf16expertsmonolithic", "vllm_unquantizedfusedmoe_triton_tritonexperts"]} - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "moe", moe_dtype: "fp8", backends: ["flashinfer_cutlass", "flashinfer_trtllm_gen"], kernel_sources: ["vllm_modeloptfp8moe_flashinfer_cutlass_flashinferexperts", "vllm_modeloptfp8moe_flashinfer_trtllm_trtllmfp8expertsmonolithic"]} - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "moe", moe_dtype: "fp8_block", backends: ["flashinfer_trtllm_gen", "triton_fused_moe"], kernel_sources: ["vllm_fp8moe_flashinfer_trtllm_trtllmfp8expertsmodular", "vllm_fp8moe_flashinfer_trtllm_trtllmfp8expertsmonolithic", "vllm_fp8moe_triton_tritonexperts"]} - - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "moe", moe_dtype: "int4_wo", backends: ["unverified"], kernel_sources: ["vllm_compressedtensorswna16moe_marlin_marlinexperts"]} - - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "moe", moe_dtype: "nvfp4", backends: ["flashinfer_trtllm_gen", "unverified"], kernel_sources: ["vllm_compressedtensorsw4a4nvfp4moe_flashinfer_cutedsl_flashinfercutedslexperts", "vllm_compressedtensorsw4a4nvfp4moe_flashinfer_trtllm_trtllmnvfp4expertsmodular", "vllm_compressedtensorsw4a4nvfp4moe_flashinfer_trtllm_trtllmnvfp4expertsmonolithic"]} + - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "moe", moe_dtype: "int4_wo", backends: ["marlin"], kernel_sources: ["vllm_compressedtensorswna16moe_marlin_marlinexperts"]} + - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "moe", moe_dtype: "nvfp4", backends: ["flashinfer_cutedsl", "flashinfer_trtllm_gen"], kernel_sources: ["vllm_compressedtensorsw4a4nvfp4moe_flashinfer_cutedsl_flashinfercutedslexperts", "vllm_compressedtensorsw4a4nvfp4moe_flashinfer_trtllm_trtllmnvfp4expertsmodular", "vllm_compressedtensorsw4a4nvfp4moe_flashinfer_trtllm_trtllmnvfp4expertsmonolithic"]} - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "moe", moe_dtype: "w4a16_mxfp4", backends: ["flashinfer_trtllm_gen", "marlin"], kernel_sources: ["vllm_compressedtensorsw4a4mxfp4moe_marlin_marlinexperts_situ_as_silu", "vllm_gptossmxfp4moe_flashinfer_trtllm_mxfp4_bf16_trtllmmxfp4expertsmonolithic"]} - {framework: "vllm", version: "0.30.0", system: "gb300", op_name: "moe", moe_dtype: "w4a8_mxfp4_mxfp8", backends: ["flashinfer_trtllm_gen"], kernel_sources: ["vllm_mxfp4moe_flashinfer_trtllm_mxfp4_mxfp8_trtllmmxfp4expertsmodular"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "moe", moe_dtype: "bfloat16", backends: ["triton_fused_moe"], kernel_sources: ["vllm_unquantizedfusedmoe_triton_tritonexperts"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "moe", moe_dtype: "fp8", backends: ["flashinfer_cutlass"], kernel_sources: ["vllm_modeloptfp8moe_flashinfer_cutlass_flashinferexperts"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "moe", moe_dtype: "fp8_block", backends: ["flashinfer_cutlass", "triton_fused_moe"], kernel_sources: ["vllm_fp8moe_flashinfer_cutlass_flashinferexperts", "vllm_fp8moe_triton_tritonexperts"]} - - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "moe", moe_dtype: "int4_wo", backends: ["unverified"], kernel_sources: ["vllm_compressedtensorswna16moe_marlin_marlinexperts"]} + - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "moe", moe_dtype: "int4_wo", backends: ["marlin"], kernel_sources: ["vllm_compressedtensorswna16moe_marlin_marlinexperts"]} - {framework: "vllm", version: "0.30.0", system: "h100_sxm", op_name: "moe", moe_dtype: "w4a16_mxfp4", backends: ["marlin", "triton_fused_moe"], kernel_sources: ["vllm_compressedtensorsw4a4mxfp4moe_marlin_marlinexperts_situ_as_silu", "vllm_gptossmxfp4moe_triton_oaitritonmxfp4expertsmonolithic"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "moe", moe_dtype: "bfloat16", backends: ["triton_fused_moe"], kernel_sources: ["vllm_unquantizedfusedmoe_triton_tritonexperts"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "moe", moe_dtype: "fp8", backends: ["flashinfer_cutlass"], kernel_sources: ["vllm_modeloptfp8moe_flashinfer_cutlass_flashinferexperts"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "moe", moe_dtype: "fp8_block", backends: ["flashinfer_cutlass", "triton_fused_moe"], kernel_sources: ["vllm_fp8moe_flashinfer_cutlass_flashinferexperts", "vllm_fp8moe_triton_tritonexperts"]} - - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "moe", moe_dtype: "int4_wo", backends: ["unverified"], kernel_sources: ["vllm_compressedtensorswna16moe_marlin_marlinexperts"]} + - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "moe", moe_dtype: "int4_wo", backends: ["marlin"], kernel_sources: ["vllm_compressedtensorswna16moe_marlin_marlinexperts"]} - {framework: "vllm", version: "0.30.0", system: "h200_sxm", op_name: "moe", moe_dtype: "w4a16_mxfp4", backends: ["marlin", "triton_fused_moe"], kernel_sources: ["vllm_compressedtensorsw4a4mxfp4moe_marlin_marlinexperts_situ_as_silu", "vllm_gptossmxfp4moe_triton_oaitritonmxfp4expertsmonolithic"]} - {framework: "vllm", version: "0.30.0", system: "l40s", op_name: "moe", moe_dtype: "bfloat16", backends: ["triton_fused_moe"], kernel_sources: ["vllm_unquantizedfusedmoe_triton_tritonexperts"]} - {framework: "vllm", version: "0.30.0", system: "l40s", op_name: "moe", moe_dtype: "fp8", backends: ["triton_fused_moe"], kernel_sources: ["vllm_modeloptfp8moe_triton_tritonexperts"]} - {framework: "vllm", version: "0.30.0", system: "l40s", op_name: "moe", moe_dtype: "fp8_block", backends: ["triton_fused_moe"], kernel_sources: ["vllm_fp8moe_triton_tritonexperts"]} - - {framework: "vllm", version: "0.30.0", system: "l40s", op_name: "moe", moe_dtype: "int4_wo", backends: ["unverified"], kernel_sources: ["vllm_compressedtensorswna16moe_marlin_marlinexperts"]} + - {framework: "vllm", version: "0.30.0", system: "l40s", op_name: "moe", moe_dtype: "int4_wo", backends: ["marlin"], kernel_sources: ["vllm_compressedtensorswna16moe_marlin_marlinexperts"]} - {framework: "vllm", version: "0.30.0", system: "l40s", op_name: "moe", moe_dtype: "w4a16_mxfp4", backends: ["marlin"], kernel_sources: ["vllm_compressedtensorsw4a4mxfp4moe_marlin_marlinexperts_situ_as_silu", "vllm_gptossmxfp4moe_marlin_marlinexperts"]} - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "moe", moe_dtype: "bfloat16", backends: ["flashinfer_cutlass", "triton_fused_moe"], kernel_sources: ["vllm_unquantizedfusedmoe_flashinfer_cutlass_flashinferexperts", "vllm_unquantizedfusedmoe_triton_tritonexperts"]} - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "moe", moe_dtype: "fp8", backends: ["flashinfer_cutlass"], kernel_sources: ["vllm_modeloptfp8moe_flashinfer_cutlass_flashinferexperts"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "moe", moe_dtype: "fp8_block", backends: ["triton_fused_moe", "unverified"], kernel_sources: ["vllm_fp8moe_deepgemm_deepgemmexperts", "vllm_fp8moe_triton_tritonexperts"]} - - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "moe", moe_dtype: "int4_wo", backends: ["unverified"], kernel_sources: ["vllm_compressedtensorswna16moe_marlin_marlinexperts"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "moe", moe_dtype: "fp8_block", backends: ["deepgemm", "triton_fused_moe"], kernel_sources: ["vllm_fp8moe_deepgemm_deepgemmexperts", "vllm_fp8moe_triton_tritonexperts"]} + - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "moe", moe_dtype: "int4_wo", backends: ["marlin"], kernel_sources: ["vllm_compressedtensorswna16moe_marlin_marlinexperts"]} - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "moe", moe_dtype: "nvfp4", backends: ["flashinfer_cutlass"], kernel_sources: ["vllm_compressedtensorsw4a4nvfp4moe_flashinfer_cutlass_flashinferexperts"]} - {framework: "vllm", version: "0.30.0", system: "rtx_pro_6000_server", op_name: "moe", moe_dtype: "w4a16_mxfp4", backends: ["marlin"], kernel_sources: ["vllm_compressedtensorsw4a4mxfp4moe_marlin_marlinexperts_situ_as_silu", "vllm_gptossmxfp4moe_marlin_marlinexperts"]} - {framework: "vllm", version: "0.31.0", system: "b300_sxm", op_name: "moe", moe_dtype: "fp8_block", backends: ["flashinfer_trtllm_gen", "triton_fused_moe"], kernel_sources: ["vllm_fp8moe_flashinfer_trtllm_trtllmfp8expertsmonolithic", "vllm_fp8moe_triton_tritonexperts"]} diff --git a/python/aisimulate/collector/registry_types.py b/python/aisimulate/collector/registry_types.py index bdeace7149..99d57cd2f3 100644 --- a/python/aisimulate/collector/registry_types.py +++ b/python/aisimulate/collector/registry_types.py @@ -7,6 +7,7 @@ from dataclasses import dataclass, field from enum import Enum +from typing import Any class PerfFile(str, Enum): @@ -35,6 +36,12 @@ def __str__(self) -> str: KDA = "kda_perf.txt" MAMBA2 = "mamba2_perf.txt" COMPUTESCALE = "computescale_perf.txt" + # computescale's second table (dynamic-vs-static fp8 quant matrix): both + # backends have always written it and the SDK consumes + # scale_matrix_perf.parquet (aisimulate_core/sdk/common.py) — the enum + # entry was simply missing, which the fail-closed finalize rejected on + # the first full vllm computescale run (2026-09-21). + SCALE_MATRIX = "scale_matrix_perf.txt" WIDEEP_CONTEXT_MLA = "wideep_context_mla_perf.txt" WIDEEP_GENERATION_MLA = "wideep_generation_mla_perf.txt" WIDEEP_CONTEXT_MOE = "wideep_context_moe_perf.txt" @@ -51,10 +58,11 @@ def __str__(self) -> str: DSA_GENERATION_MODULE = "dsa_generation_module_perf.txt" # GLM-5.2 shares one topk index across `index_topk_freq` layers: only 1 # layer per group computes the indexer (mqa+topk+index-K store), the rest - # reuse it (skip_indexer). These names are worker selectors; skip rows - # share the canonical DSA phase tables above, distinguished by op_name. - DSA_CONTEXT_MODULE_SKIP_INDEXER = "dsa_context_module_skip_indexer_perf.txt" - DSA_GENERATION_MODULE_SKIP_INDEXER = "dsa_generation_module_skip_indexer_perf.txt" + # reuse it (skip_indexer). The skip-layer rows live in the SAME two DSA + # tables above, tagged by an op_name "_skip_indexer" suffix (the SDK reads + # them from there: sdk/database test_dsa_skip_indexer_fail_open), so the + # skip ops declare DSA_*_MODULE and select the variant through + # OpEntry.run_kwargs. per_layer = (1/freq)*full + (1-1/freq)*skip. # MiniMax-M3 MSA (block-sparse GQA) full-module data — same row schema as # the DSA module files, keyed by architecture. MSA_CONTEXT_MODULE = "msa_context_module_perf.txt" @@ -67,11 +75,12 @@ def __str__(self) -> str: DSV4_HCA_CONTEXT_MODULE = "dsv4_hca_context_module_perf.txt" DSV4_CSA_GENERATION_MODULE = "dsv4_csa_generation_module_perf.txt" DSV4_HCA_GENERATION_MODULE = "dsv4_hca_generation_module_perf.txt" - # DeepSeek-V4 sparse-kernel data — bench-collected (paged_mqa_logits + - # hca_attn + csa_attn), each 1:1 with its owning CSA/HCA module rows. + # DeepSeek-V4 sparse-kernel data — the CSA indexer scoring sub-kernel, + # bench-collected 1:1 with its owning CSA module rows. The hca_attn / + # csa_attn FMLA sidecars were removed end to end on 2026-10-15: their SDK + # readers had been retired (no query path; csa_attn never shipped data) + # and the registry-only ops failed every shard that named them. DSV4_PAGED_MQA_LOGITS_MODULE = "dsv4_paged_mqa_logits_module_perf.txt" - DSV4_HCA_ATTN_MODULE = "dsv4_hca_attn_module_perf.txt" - DSV4_CSA_ATTN_MODULE = "dsv4_csa_attn_module_perf.txt" # DeepSeek-V4 CSA topk_512 degenerate-vs-representative DELTA calibration. # SGLang 0.5.14 rows qualify flat/top_last by the executed v1/v2 variant. DSV4_CSA_TOPK_CALIB = "dsv4_csa_topk_calib_perf.txt" @@ -104,9 +113,10 @@ class OpEntry: Exactly one of ``module`` (unversioned) or ``versions`` (versioned) must be provided. This invariant is validated at construction time. - ``perf_filename`` is the physical output table. ``worker_perf_filename`` - optionally overrides only the basename passed to the worker's legacy - ``perf_filename`` dispatch argument, preserving the output directory. + ``perf_filename`` is the physical output table. Ops that share one + run_func and one table in different modes select the mode with + ``run_kwargs`` (sglang dsa_*_module_skip_indexer); multi-table producers + name their extra tables in ``extra_perf_filenames``. """ op: str @@ -123,10 +133,46 @@ class OpEntry: # unverified_sms=(120,) — debugged elsewhere, not validated on these SMs unverified: bool = False unverified_sms: tuple[int, ...] = () - worker_perf_filename: str | None = None + # Further tables the SAME run_func writes from one measurement (passed to + # it as ``extra_perf_filenames``); finalize binds every one of them to + # this producer's checkpoint. compute_scale is the only user so far + # (computescale_perf + scale_matrix_perf; owner decision 2026-09-28). + extra_perf_filenames: tuple[str, ...] = () + # Fixed keyword arguments the executor binds to run_func (functools.partial, + # next to perf_filename / extra_perf_filenames). This is how one run_func + # serves several registry ops that write the SAME table in different modes: + # sglang dsa_*_module_skip_indexer passes skip_indexer=True and shares the + # dsa_*_module table with the full-module op. Before this field the worker + # inferred the mode from a "skip_indexer" perf_filename that no collector + # ever wrote, and --resume finalization (which trusts the declared table) + # raised "open checkpoint event has no regular staging table" on every + # skip_indexer shard (b200_sxm sglang 0.5.21, GitLab job 469988017, + # 2026-10-05). A tuple of (name, value) pairs keeps the entry hashable. + run_kwargs: tuple[tuple[str, Any], ...] = () def __post_init__(self) -> None: if not self.module and not self.versions: raise ValueError(f"OpEntry '{self.op}': must specify 'module' or 'versions'") if self.module and self.versions: raise ValueError(f"OpEntry '{self.op}': cannot specify both 'module' and 'versions'") + names = [str(name) for name in self.perf_filenames] + if len(names) != len(set(names)): + raise ValueError(f"OpEntry '{self.op}': duplicate perf tables {names}") + keys = [pair[0] for pair in self.run_kwargs] + if any(not isinstance(key, str) or not key for key in keys) or len(keys) != len(set(keys)): + raise ValueError(f"OpEntry '{self.op}': run_kwargs must be unique non-empty names, got {keys}") + reserved = {"perf_filename", "extra_perf_filenames", "device"} + if reserved & set(keys): + raise ValueError( + f"OpEntry '{self.op}': run_kwargs may not rebind executor-owned {sorted(reserved & set(keys))}" + ) + + @property + def run_kwargs_dict(self) -> dict[str, Any]: + """The fixed keyword arguments as a dict (what the executor binds).""" + return dict(self.run_kwargs) + + @property + def perf_filenames(self) -> tuple[str, ...]: + """Every table this producer writes, primary first.""" + return (self.perf_filename, *self.extra_perf_filenames) diff --git a/python/aisimulate/collector/sglang/collect_attn.py b/python/aisimulate/collector/sglang/collect_attn.py index bab7c862ad..f8787758ed 100644 --- a/python/aisimulate/collector/sglang/collect_attn.py +++ b/python/aisimulate/collector/sglang/collect_attn.py @@ -9,14 +9,17 @@ KV-cache setup, backend dispatch, and perf logging for the SGLang runtime. """ -__compat__ = "sglang==0.5.14" +# 0.5.21 since 2026-10-01 (H20/sm90 collector port: op_smoke + path gates in the v0.5.21 image; +# findings hopper_sglang_collector_port_0514_to_0521_2026_10_01). +# One release per file (owner policy 2026-09-20); earlier code is `git log -- ` away. +__compat__ = "sglang==0.5.21" import math import os +from importlib.metadata import version as _dist_version # setuptools/pkg_resources is absent from sglang>=0.5.21 images from types import SimpleNamespace from typing import NamedTuple -import pkg_resources import torch from collector.case_generator import ( get_attention_context_shape_sweeps, @@ -41,6 +44,9 @@ class Timing(NamedTuple): # Mock objects to satisfy RadixAttention dependencies class MockModelConfig: + def get_max_num_attention_heads(self) -> int: # model_config.py:1504 @0.5.21 (triton_backend.py:255) + return int(self.num_attention_heads) + def __init__( self, num_attention_heads, @@ -68,6 +74,11 @@ def __init__( self.is_multimodal = False self.hidden_size = num_attention_heads * head_dim self.is_local_attention_model = attention_chunk_size is not None + # sglang 0.5.16: configs/hybrid_arch.py:30-31 reads + # model_config.linear_attn_registry_result (ModelConfig property, + # model_config.py:707) to detect hybrid linear-attention archs; the + # dense-attention mock has none. + self.linear_attn_registry_result = None class MockHFConfig: def __init__(self, *, num_attention_heads, num_key_value_heads, head_dim, v_head_dim, architecture): @@ -81,6 +92,13 @@ def __init__(self, *, num_attention_heads, num_key_value_heads, head_dim, v_head self.hidden_size = num_attention_heads * head_dim self.attn_logit_softcapping = None + def get_text_config(self): + # HF PretrainedConfig API: sglang 0.5.16 unwraps multimodal + # configs through it (configs/hybrid_arch.py:49, + # linear_attn_model_registry.py:56); a text-only mock is its own + # text config. + return self + self.hf_config = MockHFConfig( num_attention_heads=num_attention_heads, num_key_value_heads=num_key_value_heads, @@ -91,7 +109,7 @@ def __init__(self, *, num_attention_heads, num_key_value_heads, head_dim, v_head self.hf_text_config = self.hf_config self.dtype = torch.bfloat16 - def get_num_kv_heads(self, tp_size): + def get_num_kv_heads(self, tp_size, dcp_size=1): # model_config.py:1512 @0.5.21 adds dcp_size (DCP replicates KV) return self.num_key_value_heads // tp_size @@ -121,9 +139,15 @@ def __init__(self, page_size: int): self.enable_two_batch_overlap = False self.disable_attn_tp_gather = False self.moe_dense_tp_size = None + # sglang 0.5.16: flashattention_backend.py:51 reads + # server_args.enable_prefill_cp (prefill context parallelism) — off + self.enable_prefill_cp = False class MockModelRunner: + def decode_num_tokens_per_req(self, *, num_draft_tokens=None) -> int: + return 1 + def __init__( self, device, @@ -136,6 +160,7 @@ def __init__( architecture=None, runtime_window_size=-1, attention_chunk_size=None, + attn_backend_name=None, ): self.device = device self.req_to_token_pool = None @@ -143,7 +168,56 @@ def __init__( self.attn_backend = None self.server_args = MockServerArgs(page_size=page_size) self.attn_cp_size = 1 # Context parallelism size; required by FlashAttentionBackend in sglang >=0.5.10 + # sglang 0.5.21 TRTLLMHAAttnBackend (the dense backend sglang picks on SM100/103) reads + # model_runner.prefill_attention_backend_str / decode_attention_backend_str to decide its + # prefill/decode split; fa3/flashinfer on Hopper never asked (B200 smoke 2026-10-04: + # AttributeError 'MockModelRunner' object has no attribute 'prefill_attention_backend_str'). + self.prefill_attention_backend_str = attn_backend_name + self.decode_attention_backend_str = attn_backend_name + # trtllm_mha_backend.py:310@0.5.21 scans model_runner.model.modules() for ENCODER_ONLY attention + # layers (expanded TARGET_VERIFY metadata); the collector has no model — an empty Module + # declares none, which is the dense-decoder answer. + self.model = torch.nn.Module() + # TRTLLMHAAttnBackend also sizes its workspaces from model_runner.max_running_requests + # (512 = the generator's --max-running-requests default). + self.max_running_requests = 512 + # sglang 0.5.16 reads the parallel geometry from ``model_runner.ps`` + # (flashattention_backend.py:183 attn_cp_size, :271-274 tp_size; the + # runner sets it at model_runner.py:262 from a ParallelState). Provide + # the real trivial ParallelState when the wrapper exists, else a + # namespace with the same fields (parallel_state_wrapper.py:6-24). + try: + from sglang.srt.distributed.parallel_state_wrapper import ParallelState + + self.ps = ParallelState.trivial(gpu_id=0) + except ImportError: + from types import SimpleNamespace + + self.ps = SimpleNamespace( + tp_rank=0, + tp_size=1, + pp_rank=0, + pp_size=1, + dp_rank=0, + dp_size=1, + attn_tp_rank=0, + attn_tp_size=1, + attn_cp_rank=0, + attn_cp_size=1, + attn_dp_rank=0, + attn_dp_size=1, + moe_ep_rank=0, + moe_ep_size=1, + moe_dp_rank=0, + moe_dp_size=1, + dcp_size=1, + gpu_id=0, + ) self.is_draft_worker = False + # sglang>=0.5.21 backends ask the runner for logits rows per decode slot + # (model_runner.py:796 decode_num_tokens_per_req; triton_backend.py:218,222); + # the kernel collectors never run speculative decoding -> 1 + self.spec_algorithm = None self.model_is_mrope = False self.sliding_window_size = attention_chunk_size if self.sliding_window_size is None and runtime_window_size >= 0: @@ -161,6 +235,11 @@ def __init__( attention_chunk_size, ) self.kv_cache_dtype = kv_cache_dtype # Default + # sglang>=0.5.21 FlashAttentionBackend reads the server-args string next to the + # torch dtype (flashattention_backend.py:192-194 kv_cache_dtype_str, drives the + # mxfp8 branch); serving sets it from server_args.kv_cache_dtype + # (model_runner.py:1481). The collector's kv_cache_dtype is that same string. + self.kv_cache_dtype_str = kv_cache_dtype self.page_size = page_size self.tp_size = 1 self.is_hybrid = False @@ -354,6 +433,72 @@ def get_generation_attention_test_cases(): return test_cases +def default_attention_backend(sm_version: int, has_attention_sink: bool) -> str | None: + """SGLang's default dense-attention backend per SM (the table run_attention_torch documents + above its call site); None for SMs outside the supported set.""" + return { + 89: "triton" if has_attention_sink else "flashinfer", + 90: "fa3", + 100: "trtllm_mha", + 103: "trtllm_mha", + 120: "triton" if has_attention_sink else "flashinfer", + }.get(sm_version) + + +_FLASHINFER_WORKSPACE_MIN_BYTES = 4 << 30 + + +def _ensure_flashinfer_workspace() -> int | None: + """Raise SGLang's FlashInfer workspace floor BEFORE the first backend of this + process is built, through the knob the pinned version actually reads. + + 0.5.21 sizes the shared decode/prefill plan buffer from + ``sglang.srt.environ.envs.SGLANG_FLASHINFER_WORKSPACE_SIZE`` (EnvInt, default + 384 MiB) and allocates it ONCE per process via ``get_buffer("flashinfer_workspace")`` + (flashinfer_backend.py:430-436); serving itself bumps the knob to 512 MiB for + Qwen2/Qwen3/MiMo architectures (:394-404). The decode sweep's wide cells need + more than the default — l40s 0.5.21: "Buffer overflow when allocating + batch_prefill_tmp_v with size 536870912 ... only 402653184 bytes available", + 671 attention_generation cases (heads 128 / kv 1 / head_dim 128, batch 32 x isl + 255..4095). ed9109ec wrote ``global_config.flashinfer_workspace_size``, an + attribute 0.5.21 no longer has — a no-op (pipeline 72343559 reproduced all 671). + Older series keep the global_config path as a fallback. Returns the floor + applied (bytes) or None when no knob was found. + + Floor history: 1 GiB (5c0f54a5) cleared 447 of the 671 l40s failures (pipeline + 72369799, job 475655072); the remaining 224 — all heads 128 / kv 1 / head_dim 128, + bf16 and fp8 KV alike, batch x isl from 32x255 up to 4x16383 — asked the tensor-core + decode planner for batch_prefill_tmp_v of 1.5-2.22 GiB ("only 1073741824 bytes + available"). 4 GiB covers the largest observed request with margin; the buffer is one + torch.empty per worker process, never touched beyond what the plan uses, and the + planner's split decision does not depend on the buffer size, so the measured kernels + are unchanged (flashinfer prefill planner sizes tmp_v from the split count, not the + workspace). + """ + try: + from sglang.srt.environ import envs + + knob = getattr(envs, "SGLANG_FLASHINFER_WORKSPACE_SIZE", None) + if knob is not None: + current = int(knob.get()) + if current < _FLASHINFER_WORKSPACE_MIN_BYTES: + knob.set(_FLASHINFER_WORKSPACE_MIN_BYTES) + return max(current, _FLASHINFER_WORKSPACE_MIN_BYTES) + except Exception: + pass + try: + from sglang.global_config import global_config + + current = getattr(global_config, "flashinfer_workspace_size", None) + if isinstance(current, int): + if current < _FLASHINFER_WORKSPACE_MIN_BYTES: + global_config.flashinfer_workspace_size = _FLASHINFER_WORKSPACE_MIN_BYTES + return max(current, _FLASHINFER_WORKSPACE_MIN_BYTES) + except Exception: + pass + return None + + @get_parallel().override( attn_tp_size=1, attn_tp_rank=0, @@ -384,6 +529,10 @@ def run_attention_torch( device="cuda:0", page_size: int | None = None, ): + from collector.sglang.runtime_compat import ensure_offline_runtime_published + + ensure_offline_runtime_published() # sglang>=0.5.20 backends read get_exec()/get_parallel() + _ensure_flashinfer_workspace() if use_fp8_context_fmha: assert use_fp8_kv_cache, "If you want to use fp8 context fmha, kv cache must be fp8" kvtype = torch.float8_e4m3fn if use_fp8_kv_cache else torch.bfloat16 @@ -415,18 +564,48 @@ def run_attention_torch( # sglang_backends map as an explicit attn_backend_name, never through # this default table. SM80/86 are outside the # supported platform set {89, 90, 100, 103, 120} and fail closed below. - attn_backend_name = { - 89: "triton" if has_attention_sink else "flashinfer", - 90: "fa3", - 100: "trtllm_mha", - 103: "trtllm_mha", - 120: "triton" if has_attention_sink else "flashinfer", - }.get(sm_version) + attn_backend_name = default_attention_backend(sm_version, has_attention_sink) if attn_backend_name is None: raise ValueError(f"No SGLang 0.5.14 attention backend mapping for SM{sm_version}") if page_size is None: page_size = 64 if attn_backend_name == "trtllm_mha" else 1 + # Label prefill rows by the compute dtype the backend actually uses in + # SGLang 0.5.21 — the input dtype is an API contract, not the compute + # dtype: + # - fa3 (SM90): whenever the KV cache is FP8 and head_dim <= 256 it + # casts Q to the KV dtype itself (flashattention_backend.py:1322-1332), + # so FP8-KV prefill IS FP8 FMHA compute. + # - trtllm_mha (TRTLLM-GEN on SM100/103): takes BF16 inputs and casts Q + # to FP8 itself when the KV cache is FP8 (trtllm_mha_backend.py:1600-1609). + # Feed BF16 and let the backend quantize; an external FP8 cast would + # break its input contract without changing the (already FP8) compute. + # - flashinfer: the prefill wrapper plans BF16 Q against the FP8 KV cache + # (flashinfer_backend.py:1833-1834) and reads it with descales + # (:1388-1389); there is no FP8 prefill compute path, so an fp8-labeled + # row would be a lie. + # Consequence: on fa3/trtllm_mha a "BF16-compute prefill on FP8 KV" + # combo does not exist — the backend force-quantizes Q — so collecting + # it would record FP8 compute under a bfloat16 label. Both refusals run + # here, before the model runner and tensors are built; the cases stay + # queued and land in the failure log as kernel_limit. + if is_context_phase: + if use_fp8_context_fmha and attn_backend_name == "flashinfer": + raise ValueError( + "FIXME(kernel-limit): SGLang flashinfer has no FP8 prefill compute path " + "(flashinfer_backend.py:1833-1834 @0.5.21 plans BF16 Q against the FP8 KV cache)" + ) + if ( + not use_fp8_context_fmha + and use_fp8_kv_cache + and (attn_backend_name == "trtllm_mha" or (attn_backend_name == "fa3" and head_dim <= 256)) + ): + raise ValueError( + f"FIXME(kernel-limit): SGLang {attn_backend_name} quantizes Q to FP8 internally when the KV cache " + "is FP8, so a BF16-compute prefill on an FP8 KV cache does not exist; this " + "combination is the fp8_context_fmha case" + ) + model_runner = MockModelRunner( torch_device, kv_cache_dtype="fp8_e4m3" if use_fp8_kv_cache else "auto", @@ -438,8 +617,15 @@ def run_attention_torch( architecture=architecture, runtime_window_size=runtime_window_size, attention_chunk_size=attention_chunk_size, + attn_backend_name=attn_backend_name, ) model_runner.kv_cache_dtype = kvtype + # sglang>=0.5.21 backends ask the runner which backend serves each phase + # (model_runner.py:1020-1021 resolved.prefill/decode; flashinfer_backend.py:331-332 + # gate the kv-access checks on it). One lane = one backend for both phases. Hopper's + # fa3 / Blackwell's trtllm_mha never read it, so only the sm89/sm120 flashinfer lane failed. + model_runner.prefill_attention_backend_str = attn_backend_name + model_runner.decode_attention_backend_str = attn_backend_name total_len = input_len if is_context_phase else input_len + 1 # TRTLLM MHA sizes its page table from context_len. @@ -474,10 +660,14 @@ def run_attention_torch( page_size=model_runner.page_size, get_kvcache=lambda: model_runner.token_to_kv_pool, ) + from collector.sglang.runtime_compat import attach_kv_index_translator + + attach_kv_index_translator(model_runner) if attn_backend_name == "flashinfer": from sglang.srt.layers.attention.flashinfer_backend import FlashInferAttnBackend + _ensure_flashinfer_workspace() # the shared workspace is allocated by the first backend of the process attn_backend = FlashInferAttnBackend(model_runner) elif attn_backend_name == "trtllm_mha": from sglang.srt.layers.attention.trtllm_mha_backend import TRTLLMHAAttnBackend @@ -621,38 +811,12 @@ def run_attention_torch( with forward_context(ForwardContext(attn_backend=attn_backend)): attn_backend.init_forward_metadata(forward_batch) - # Label prefill rows by the compute dtype the backend actually uses in - # SGLang 0.5.14 — the input dtype is an API contract, not the compute - # dtype: - # - fa3 (SM90): whenever the KV cache is FP8 and head_dim <= 256 it - # casts Q to the KV dtype itself (flashattention_backend.py:857-872), - # so FP8-KV prefill IS FP8 FMHA compute. - # - trtllm_mha (TRTLLM-GEN on SM100/103): requires BF16 inputs and - # quantizes Q internally via scaled_fp8_quant when the KV cache is - # FP8 (trtllm_mha_backend.py:154-158, 291-301). Feed BF16 and let - # the backend quantize; an external FP8 cast would break its input - # contract without changing the (already FP8) compute. - # - flashinfer: BF16 Q reads the FP8 KV cache with descales; there is - # no FP8 prefill compute path, so an fp8-labeled row would be a lie. - # Consequence: on fa3/trtllm_mha a "BF16-compute prefill on FP8 KV" - # combo does not exist — the backend force-quantizes Q — so collecting - # it would record FP8 compute under a bfloat16 label. Fail closed. - if is_context_phase: - if use_fp8_context_fmha: - if attn_backend_name == "flashinfer": - raise ValueError("SGLang 0.5.14 flashinfer has no FP8 prefill compute path") - if attn_backend_name != "trtllm_mha": - q = q.to(kvtype) - k = k.to(kvtype) - v = v.to(kvtype) - elif use_fp8_kv_cache and ( - attn_backend_name == "trtllm_mha" or (attn_backend_name == "fa3" and head_dim <= 256) - ): - raise ValueError( - f"SGLang 0.5.14 {attn_backend_name} quantizes Q to FP8 internally when the KV cache " - "is FP8, so a BF16-compute prefill on an FP8 KV cache does not exist; this " - "combination is the fp8_context_fmha case" - ) + # fp8_context_fmha (labels: see the refusals before the model runner): + # trtllm_mha casts Q itself; every other lane takes FP8 inputs. + if is_context_phase and use_fp8_context_fmha and attn_backend_name != "trtllm_mha": + q = q.to(kvtype) + k = k.to(kvtype) + v = v.to(kvtype) # Mirror the serving call contract: only sink-carrying models pass the # ``sinks`` kwarg (e.g. gpt_oss.py), and SGLang 0.5.14 @@ -707,7 +871,7 @@ def run_iter(): } ], framework="SGLang", - version=pkg_resources.get_distribution("sglang").version, + version=_dist_version("sglang"), device_name=torch.cuda.get_device_name(device), op_name=op_name, kernel_source=attn_backend_name, diff --git a/python/aisimulate/collector/sglang/collect_attn_encoder.py b/python/aisimulate/collector/sglang/collect_attn_encoder.py index 81be051301..d9a7bd8e0c 100644 --- a/python/aisimulate/collector/sglang/collect_attn_encoder.py +++ b/python/aisimulate/collector/sglang/collect_attn_encoder.py @@ -5,23 +5,43 @@ SM dispatch mirrors ``VisionAttention._determine_attention_backend``: -- CUDA SM == 90 (Hopper) -> ``flash_attn_varlen_func`` (FA3) -- CUDA SM == 100 (Blackwell) -> ``flash_attn_varlen_func(ver=4)`` (FA4) -- other CUDA (SM<90, SM120) -> ``context_attention_fwd`` (Triton) +- capability major 9 (SM90, Hopper) -> ``flash_attn_varlen_func`` (FA3) +- capability major 10 (SM100/SM103, B200/B300) -> ``flash_attn_varlen_func(ver=4)`` (FA4) +- other CUDA (SM<90, SM120) -> ``context_attention_fwd`` (Triton) + +Serving branches on the capability MAJOR (layers/attention/vision.py:1298-1303 @0.5.21: +``major == 9 -> fa3``, ``major == 10 -> fa4``, else ``triton_attn``), so SM103 shares the FA4 +lane with SM100. The collector's earlier ``sm == 100`` test sent SM103 down the Triton path +(B300 encoder gate failed, 2026-10-05) — a wrong-kernel measurement, not a crash. Quant: bf16 only. SGLang upstream does not support fp8 ViT FMHA. """ -__compat__ = "sglang==0.5.14" +# 0.5.21 since 2026-10-01 (H20/sm90 collector port: op_smoke + path gates in the v0.5.21 image; +# findings hopper_sglang_collector_port_0514_to_0521_2026_10_01). +# One release per file (owner policy 2026-09-20); earlier code is `git log -- ` away. +__compat__ = "sglang==0.5.21" +from importlib.metadata import ( + version as _dist_version, # setuptools/pkg_resources is absent from recent framework images +) from typing import NamedTuple -import pkg_resources import torch from collector.case_generator import get_attention_encoder_head_configs, get_attention_encoder_shape_sweeps from collector.helper import benchmark_with_power, get_sm_version, log_perf +def serving_vision_backend(sm_version: int) -> str: + """The vision-attention backend sglang selects for a CUDA capability (vision.py:1298-1303 @0.5.21).""" + major = sm_version // 10 + if major == 9: + return "fa3" + if major == 10: + return "fa4" + return "triton_attn" + + class Timing(NamedTuple): mean: float @@ -80,11 +100,15 @@ def _build_kernel_runner( max_seqlen = seq_len softmax_scale = head_dim**-0.5 - sm = get_sm_version() # 90 / 100 / 120 / ... + sm = get_sm_version() # 90 / 100 / 103 / 120 / ... - if sm == 90: + backend = serving_vision_backend(sm) + if backend == "fa3": # Matches VisionFlash3Attention.forward. - from sglang.jit_kernel.flash_attention import flash_attn_varlen_func + try: # sglang>=0.5.21 (layers/attention/vision.py:55) + from sglang.kernels.ops.attention.flash_attention import flash_attn_varlen_func + except ImportError: + from sglang.jit_kernel.flash_attention import flash_attn_varlen_func def run_iter(): flash_attn_varlen_func( @@ -101,9 +125,12 @@ def run_iter(): return run_iter, "flash_attention_v3" - if sm == 100: + if backend == "fa4": # Matches VisionFlash4Attention.forward. - from sglang.jit_kernel.flash_attention import flash_attn_varlen_func + try: # sglang>=0.5.21 (layers/attention/vision.py:55) + from sglang.kernels.ops.attention.flash_attention import flash_attn_varlen_func + except ImportError: + from sglang.jit_kernel.flash_attention import flash_attn_varlen_func def run_iter(): flash_attn_varlen_func( @@ -120,10 +147,13 @@ def run_iter(): return run_iter, "flash_attention_v4" - # SM<90 or SM>100: Triton path matching VisionTritonAttention.forward. - from sglang.srt.layers.attention.triton_ops.prefill_attention import ( - context_attention_fwd, - ) + # every other capability (SM<90, SM120): Triton path matching VisionTritonAttention.forward. + try: # sglang>=0.5.21 (layers/attention/vision.py imports it from kernels/ops/attention) + from sglang.kernels.ops.attention.prefill_attention import context_attention_fwd + except ImportError: + from sglang.srt.layers.attention.triton_ops.prefill_attention import ( + context_attention_fwd, + ) seq_lens = torch.full( (batch_size,), @@ -158,6 +188,9 @@ def run_encoder_attention_torch( perf_filename, device="cuda:0", ): + from collector.sglang.runtime_compat import ensure_offline_runtime_published + + ensure_offline_runtime_published() # sglang>=0.5.20 backends read get_exec()/get_parallel() torch_device = torch.device(device) torch.cuda.set_device(device) @@ -193,7 +226,7 @@ def run_encoder_attention_torch( } ], framework="SGLang", - version=pkg_resources.get_distribution("sglang").version, + version=_dist_version("sglang"), device_name=torch.cuda.get_device_name(device), op_name="encoder_attention", kernel_source=backend_tag, diff --git a/python/aisimulate/collector/sglang/collect_computescale.py b/python/aisimulate/collector/sglang/collect_computescale.py index 81a1c1dbce..5325e21e99 100644 --- a/python/aisimulate/collector/sglang/collect_computescale.py +++ b/python/aisimulate/collector/sglang/collect_computescale.py @@ -1,16 +1,29 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Measure SGLang FP8 activation quantization overhead for static-FP8 GEMM.""" +"""Measure SGLang FP8 activation quantization overhead for static-FP8 GEMM. -__compat__ = "sglang==0.5.14" +One measurement, two tables — the registry entry declares the second one as +``extra_perf_filenames`` and the executor hands both names in, so finalize +binds both to this producer's checkpoint: + perf_filename -> computescale_perf.txt (dynamic per-token quant minus static quant) + extra_perf_filenames[0] -> scale_matrix_perf.txt (the static per-tensor quant alone) +""" + +# 0.5.21 since 2026-10-01 (H20/sm90 collector port: op_smoke + path gates in the v0.5.21 image; +# findings hopper_sglang_collector_port_0514_to_0521_2026_10_01). +# One release per file (owner policy 2026-09-20); earlier code is `git log -- ` away. +__compat__ = "sglang==0.5.21" + +from importlib.metadata import version as _dist_version # setuptools/pkg_resources is absent from sglang>=0.5.21 images -import pkg_resources import torch from collector.case_generator import get_compute_scale_case_specs from collector.helper import benchmark_with_power, get_sm_version, log_perf from sgl_kernel import sgl_per_token_quant_fp8 +_OUTSIDE_LOOP_COUNT = 5 + def get_computescale_test_cases(): if get_sm_version() <= 86: @@ -24,37 +37,48 @@ def _static_quantize_e4m3_per_tensor(x: torch.Tensor, scale: torch.Tensor, out: return out -def run_computescale(m, k, *, perf_filename, device="cuda:0"): +def _setup(m, k, device): device = torch.device(device) torch.cuda.set_device(device) torch.set_default_device(device) - x = torch.randn((m, k), dtype=torch.bfloat16, device=device) + return device, x + + +def _bench_dynamic(device, x): + m = x.shape[0] dynamic_out = torch.empty_like(x, dtype=torch.float8_e4m3fn) dynamic_scale = torch.empty((m, 1), dtype=torch.float32, device=device) - static_out = torch.empty_like(x, dtype=torch.float8_e4m3fn) - static_scale = torch.tensor([1.0], dtype=torch.float32, device=device) - outside_loop_count = 5 - def dynamic_kernel_func(): - for _ in range(outside_loop_count): + def kernel_func(): + for _ in range(_OUTSIDE_LOOP_COUNT): sgl_per_token_quant_fp8(x, dynamic_out, dynamic_scale) - with benchmark_with_power(device=device, kernel_func=dynamic_kernel_func, repeat_n=1) as dynamic_results: + with benchmark_with_power(device=device, kernel_func=kernel_func, repeat_n=1) as results: pass + return results["latency_ms"] / _OUTSIDE_LOOP_COUNT, results["power_stats"] + - dynamic_latency = dynamic_results["latency_ms"] / outside_loop_count +def _bench_static(device, x): + static_out = torch.empty_like(x, dtype=torch.float8_e4m3fn) + static_scale = torch.tensor([1.0], dtype=torch.float32, device=device) - def static_kernel_func(): - for _ in range(outside_loop_count): + def kernel_func(): + for _ in range(_OUTSIDE_LOOP_COUNT): _static_quantize_e4m3_per_tensor(x, static_scale, static_out) - with benchmark_with_power(device=device, kernel_func=static_kernel_func, repeat_n=1) as static_results: + with benchmark_with_power(device=device, kernel_func=kernel_func, repeat_n=1) as results: pass + return results["latency_ms"] / _OUTSIDE_LOOP_COUNT, results["power_stats"] + - static_latency = static_results["latency_ms"] / outside_loop_count +def run_computescale(m, k, *, perf_filename, extra_perf_filenames, device="cuda:0"): + (scale_matrix_filename,) = extra_perf_filenames + device, x = _setup(m, k, device) + dynamic_latency, dynamic_power = _bench_dynamic(device, x) + static_latency, static_power = _bench_static(device, x) compute_scale_latency = max(0.0, dynamic_latency - static_latency) - version = pkg_resources.get_distribution("sglang").version + version = _dist_version("sglang") if not log_perf( item_list=[{"m": m, "k": k, "quant_dtype": "fp8", "latency": compute_scale_latency}], @@ -64,9 +88,10 @@ def static_kernel_func(): op_name="compute_scale", kernel_source="sglang", perf_filename=perf_filename, - power_stats=dynamic_results["power_stats"], + power_stats=dynamic_power, ): raise RuntimeError(f"Failed to persist SGLang compute scale performance row to {perf_filename}") + if not log_perf( item_list=[{"m": m, "k": k, "quant_dtype": "fp8", "latency": static_latency}], framework="SGLang", @@ -74,7 +99,7 @@ def static_kernel_func(): device_name=torch.cuda.get_device_name(device), op_name="scale_matrix", kernel_source="sglang", - perf_filename="scale_matrix_perf.txt", - power_stats=static_results["power_stats"], + perf_filename=scale_matrix_filename, + power_stats=static_power, ): - raise RuntimeError("Failed to persist SGLang scale matrix performance row to scale_matrix_perf.txt") + raise RuntimeError(f"Failed to persist SGLang scale matrix performance row to {scale_matrix_filename}") diff --git a/python/aisimulate/collector/sglang/collect_dsv4_attn.py b/python/aisimulate/collector/sglang/collect_dsv4_attn.py index 7f9f42c567..107c924a8e 100644 --- a/python/aisimulate/collector/sglang/collect_dsv4_attn.py +++ b/python/aisimulate/collector/sglang/collect_dsv4_attn.py @@ -23,10 +23,13 @@ --batch-sizes 1,4 --seq-lens 128,1024 """ -# Requires stock SGLang 0.5.14 with its matching ``sgl-kernel`` package. +# Requires stock SGLang 0.5.21 with its matching ``sgl-kernel`` package. from __future__ import annotations -__compat__ = "sglang==0.5.14" +# 0.5.21 since 2026-10-01 (H20/sm90 collector port: op_smoke + path gates in the v0.5.21 image; +# findings hopper_sglang_collector_port_0514_to_0521_2026_10_01). +# One release per file (owner policy 2026-09-20); earlier code is `git log -- ` away. +__compat__ = "sglang==0.5.21" import argparse import contextlib @@ -39,6 +42,7 @@ import subprocess import sys import tempfile +import time import traceback from collections.abc import Iterable from fractions import Fraction @@ -54,10 +58,26 @@ # few shapes which the bench's own warmup JIT-compiles on first use. os.environ["SGLANG_JIT_DEEPGEMM_PRECOMPILE"] = "0" try: - from helper import _resolve_local_model_path, benchmark_with_power, log_perf + from helper import ( + PeriodicGarbageCollector, + PhaseClock, + _resolve_local_model_path, + benchmark_with_power, + log_perf, + render_phase_totals, + worker_heartbeat, + ) except ModuleNotFoundError: sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - from helper import _resolve_local_model_path, benchmark_with_power, log_perf + from helper import ( + PeriodicGarbageCollector, + PhaseClock, + _resolve_local_model_path, + benchmark_with_power, + log_perf, + render_phase_totals, + worker_heartbeat, + ) try: from collector.sglang.runtime_limits import ( @@ -504,7 +524,16 @@ def _resolve_model_path( # model_type to "deepseek_v3" mirrors what sglang's # _load_deepseek_temp_model produces internally, so AutoConfig succeeds # and the V4 model class is still selected via the architectures field. - config["model_type"] = "deepseek_v3" + try: + # sglang>=0.5.21 ships a native DeepseekV4Config (configs/deepseek_v4.py, model_type + # "deepseek_v4") and its DSV4 model reads V4-only fields from it + # (hc_pre_from_prev_sublayer, models/deepseek_v4.py:2679); the v3 rewrite below + # would hand the model a DeepseekV3Config and fail there. + import sglang.srt.configs.deepseek_v4 # noqa: F401 + + config["model_type"] = "deepseek_v4" + except ImportError: + config["model_type"] = "deepseek_v3" # gemm_type "fp8_block" overrides disable_weight_quant: we MUST keep the # fp8 quantization_config so sglang dispatches projections to DeepGEMM. @@ -689,7 +718,11 @@ def _derive_csa_context_pool_cap( if not isinstance(configurator, DSV4PoolConfigurator): raise TypeError(f"expected SGLang 0.5.14 DSV4PoolConfigurator, got {type(configurator).__name__}") - page_size = int(model_runner.server_args.page_size) + from collector.sglang.runtime_compat import resolved_arg + + page_size = int( + resolved_arg(model_runner.server_args, "page_size") + ) # 0.5.21: derived, read through the resolving view if page_size != _DSV4_CUDA_PAGE_SIZE: raise RuntimeError(f"expected DSV4 CUDA page_size={_DSV4_CUDA_PAGE_SIZE}, got {page_size}") window_size = int(model_runner.model_config.window_size) @@ -705,7 +738,17 @@ def _derive_csa_context_pool_cap( if not chunk_eligible_shapes: raise RuntimeError(f"DSV4 CSA context pool derivation has no shape within effective_chunk={effective_chunk}") - profiled_bytes = int(model_runner._profile_available_bytes(model_runner.pre_model_load_memory)) + # sglang>=0.5.21 moved the pool profiling off ModelRunner onto the KV-cache + # configurator (mem_cache/kv_cache_configurator.py:2177 _profile_available_bytes, + # same free-memory-minus-slack arithmetic). Serving builds that object inside + # alloc_memory_pool (model_runner.py:885 init_kv_cache_configurator, then + # configure); build it the same way here, before the pool exists. + _profile = getattr(model_runner, "_profile_available_bytes", None) + if _profile is None: + if getattr(model_runner, "kv_cache_configurator", None) is None: + model_runner.init_kv_cache_configurator() + _profile = model_runner.kv_cache_configurator._profile_available_bytes + profiled_bytes = int(_profile(model_runner.pre_model_load_memory)) profiled_config = configurator.calculate_pool_sizes(profiled_bytes, page_size) compress_ratio = ATTN_KIND_TO_COMPRESS_RATIO["csa"] if page_size % compress_ratio != 0: @@ -908,20 +951,18 @@ def _load_model_runner( _set_envs_and_config(server_args) model_config = ModelConfig.from_server_args(server_args) + from collector.sglang.runtime_compat import init_runtime_config + + _runner_parallel_kwargs = init_runtime_config(server_args, gpu_id, nccl_port=nccl_port, model_config=model_config) with _tp_load_model_patch(tp_size): model_runner = ModelRunner( model_config=model_config, # Use sglang's own __post_init__-derived value, not a collector knob. mem_fraction_static=server_args.mem_fraction_static, gpu_id=gpu_id, - tp_rank=0, - tp_size=1, - pp_rank=0, - pp_size=1, - moe_ep_rank=0, - moe_ep_size=1, nccl_port=nccl_port, server_args=server_args, + **_runner_parallel_kwargs, ) derived_requirements = None if csa_context_shapes is not None: @@ -931,7 +972,18 @@ def _load_model_runner( model_runner, csa_context_shapes, ) - model_runner.server_args.max_total_tokens = derived_requirements[0] + try: + model_runner.server_args.max_total_tokens = derived_requirements[0] + except AttributeError: + # sglang>=0.5.21: ServerArgs is read-only once resolved (server_args.py:396); + # post-publish config changes go through the context's config bags + # (runtime_context.py:1080 override(source, **fields) — the same call + # serving uses for control-plane updates, tokenizer_manager.py:2204). The + # configurator reads the cap from get_schedule().max_total_tokens + # (kv_cache_configurator.py:2269 _apply_token_constraints). + from sglang.srt.runtime_context import get_context + + get_context().override("collector-dsv4-csa-context", max_total_tokens=derived_requirements[0]) # SGLang 0.5.14 separates model construction from serving-state setup. model_runner.alloc_memory_pool() if derived_requirements is not None: @@ -1036,7 +1088,11 @@ def _make_reqs( req.full_untruncated_fill_ids = array("q", req.origin_input_ids) req.fill_len = full_len req.logprob_start_len = 0 - req.set_extend_input_len(seq_len if prefix_len else full_len) + _n = seq_len if prefix_len else full_len + if hasattr(req, "set_extend_input_len"): + req.set_extend_input_len(_n) + else: # sglang 0.5.16: schedule_batch.py:1153 set_extend_range(start, end) + req.set_extend_range(full_len - _n, full_len) req.swa_evicted_seqlen = swa_evicted_seqlen if decode: req.cached_tokens = 0 @@ -1113,7 +1169,16 @@ def _build_forward_batch( req.output_ids.append(0) batch.prepare_for_decode() - forward_batch = ForwardBatch.init_new(batch, model_runner) + import inspect as _inspect + + # sglang 0.5.16: return_hidden_states_before_norm is a required keyword of + # ForwardBatch.init_new; older pins reject it (collect_msa_module.py:1015 pattern) + _fb_kw = ( + {"return_hidden_states_before_norm": False} + if "return_hidden_states_before_norm" in _inspect.signature(ForwardBatch.init_new).parameters + else {} + ) + forward_batch = ForwardBatch.init_new(batch, model_runner, **_fb_kw) model_runner.attn_backend.init_forward_metadata(forward_batch) return forward_batch @@ -1209,6 +1274,7 @@ def timed_kernel(): "used_cuda_graph": True, "power_stats": result.get("power_stats"), "throttled": bool(result.get("throttled", False)), + "phase_seconds": result.get("phase_seconds"), } @@ -1345,10 +1411,29 @@ def run_dsv4_mla_module( seq_lens_by_prefix.get(cur_prefix, default_seq_lens) if seq_lens_by_prefix is not None else default_seq_lens ) ] + phase_totals: dict[str, float] = {} + # gc.collect() after every point cost 0.2-0.4 s (half of a context sweep); collect periodically instead. + point_gc = PeriodicGarbageCollector() # SGLang derives compressed-pool sizes from max_total_tokens. Direct # probes with an explicit value must account for both the full pool and # the smaller page-rounded SWA tail pool before ModelRunner construction. + # + # HCA always takes this path since sglang 0.5.21: the ratio-128 stage prices a + # full token at ~1/128 of a paged slot (pool_configurator.py:1264-1283 + # _compressed_bytes_per_full_token / _get_bytes_per_full_token), so the + # framework-derived pool on an otherwise empty 140 GB device reaches ~2e9 full + # tokens, and the hybrid allocator's int64 full_to_swa_index_mapping + # (mem_cache/allocator/swa.py:99-108, size + page_size entries, not priced by the + # configurator) then needs ~15 GiB that the sizing already handed to the pools + # (OOM in SWATokenToKVPoolAllocator.__init__, kv_cache_configurator.py:2080, + # H20 2026-10-01). Serving never meets this corner because weights leave far + # less free memory; the collector sizes the pool to what the planned shapes + # allocate instead. sglang treats the value as an upper bound only + # (kv_cache_configurator.py:2269-2277 _apply_token_constraints: min(profiled, + # user cap)), so no cell gains capacity it would not have in serving. + if max_total_tokens is None and attn_kind == "hca": + max_total_tokens = 0 if max_total_tokens is not None: max_full_alloc = 0 max_swa_alloc = 0 @@ -1495,6 +1580,7 @@ def run_dsv4_mla_module( skipped_shapes.append((batch_size, seq_len, cur_prefix, "KVPoolSWACapacity")) continue print(f"\n{mode}: batch_size={batch_size}, seq_len={seq_len}, prefix_len={cur_prefix}") + point_clock = PhaseClock() try: forward_batch = _build_forward_batch( model_runner, @@ -1503,6 +1589,7 @@ def run_dsv4_mla_module( is_prefill=is_prefill, prefix_len=cur_prefix, ) + point_clock.mark("forward_batch") hidden_states, positions = _make_inputs( model_runner, batch_size=batch_size, @@ -1511,6 +1598,7 @@ def run_dsv4_mla_module( device=device, prefix_len=cur_prefix, ) + point_clock.mark("inputs") def kernel_func(model_runner=model_runner): from sglang.srt.model_executor.forward_context import ForwardContext, forward_context @@ -1536,6 +1624,8 @@ def kernel_func(model_runner=model_runner): graph_repeat=graph_repeat, device=device, ) + point_clock.mark("bench") + point_clock.absorb(stats.get("phase_seconds"), prefix="bench_") print( f" latency mean={stats['mean_ms']:.4f} ms, " f"median={stats['median_ms']:.4f} ms, " @@ -1593,7 +1683,7 @@ def kernel_func(model_runner=model_runner): ("req_to_token_pool.clear", model_runner.req_to_token_pool.clear), ("token_to_kv_pool_allocator.clear", model_runner.token_to_kv_pool_allocator.clear), ("torch.cuda.empty_cache", torch.cuda.empty_cache), - ("gc.collect", gc.collect), + ("gc.collect", point_gc.maybe_collect), ): try: _cleanup_step() @@ -1606,6 +1696,12 @@ def kernel_func(model_runner=model_runner): f"dsv4-flash {sweep_label} bs={batch_size} sl={seq_len} " f"prefix={cur_prefix}: cleanup failed: {'; '.join(cleanup_errors)}" ) + point_clock.mark("cleanup") + point_clock.add_to(phase_totals) + print( + f" [dsv4 point timing] {sweep_label} bs={batch_size} sl={seq_len} prefix={cur_prefix}: " + + point_clock.render() + ) finally: final_cleanup_errors = [] del model_runner @@ -1633,6 +1729,7 @@ def kernel_func(model_runner=model_runner): if skipped_shapes: skipped_str = ", ".join(f"(bs={b},sl={s},prefix={p},reason={r})" for b, s, p, r in skipped_shapes) print(f"[WARN] dsv4-flash {sweep_label}: skipped or errored shapes: {skipped_str}") + print(f"[dsv4 timing] {sweep_label}: {render_phase_totals(phase_totals)}") print(f"[dsv4-collector] {sweep_label}: {summary}") if ok_count == 0 or error_count > 0: raise RuntimeError(f"dsv4-flash {sweep_label}: {summary}") @@ -1726,11 +1823,20 @@ def _run_once(nccl_port: int) -> tuple[int, str]: stderr=subprocess.STDOUT, cwd=os.path.dirname(os.path.abspath(__file__)), ) - try: - proc.wait(timeout=3600) # up to 1 hour per (kind, tp, gemm, bs) - except subprocess.TimeoutExpired: - proc.kill() - proc.wait() + # Up to 1 hour per (kind, tp, gemm, bs). The child reports no liveness of its + # own, so beat while waiting: the executor's stall watchdog + # (STALL_TIMEOUT_SEC, 1800 s default) would otherwise fail a healthy sweep. + deadline = time.monotonic() + 3600 + while True: + try: + proc.wait(timeout=60) + break + except subprocess.TimeoutExpired: + if time.monotonic() >= deadline: + proc.kill() + proc.wait() + break + worker_heartbeat() try: with open(log_path, encoding="utf-8", errors="replace") as logf: diff --git a/python/aisimulate/collector/sglang/collect_dsv4_megamoe.py b/python/aisimulate/collector/sglang/collect_dsv4_megamoe.py index 077455c711..079b8301b6 100644 --- a/python/aisimulate/collector/sglang/collect_dsv4_megamoe.py +++ b/python/aisimulate/collector/sglang/collect_dsv4_megamoe.py @@ -446,7 +446,10 @@ def copy_pre_dispatch(hidden_states, topk_ids, topk_weights, buffer, num_tokens: try: from sglang.kernels.ops.attention.dsv4 import mega_moe_pre_dispatch except ModuleNotFoundError: - from sglang.jit_kernel.deepseek_v4 import mega_moe_pre_dispatch + try: + from sglang.jit_kernel.deepseek_v4 import mega_moe_pre_dispatch + except ImportError: # sglang>=0.5.21: kernels/ops/moe/dsv4.py:156 + from sglang.kernels.ops.moe.dsv4 import mega_moe_pre_dispatch def sglang_jit_pre_dispatch(hidden_states, topk_ids, topk_weights, buffer, num_tokens: int): del num_tokens diff --git a/python/aisimulate/collector/sglang/collect_gdn.py b/python/aisimulate/collector/sglang/collect_gdn.py index 209e5f3a33..b47c699e39 100644 --- a/python/aisimulate/collector/sglang/collect_gdn.py +++ b/python/aisimulate/collector/sglang/collect_gdn.py @@ -73,20 +73,11 @@ # unaffected) -- cosmetic relocations (signatures independently re-diffed # byte-identical at both locations), fixed with version-conditional imports # mirroring collect_gemm.py's established try/except-import pattern (this -# file had none of its own before this bump). 0.5.15/0.5.16 stay excluded: -# never verified, and version_resolver's __compat__ grammar (AND-of- -# comparators only, no OR) has no way to express a true two-point set either -# -- this bounded range with the two untested base releases carved out via -# != is the closest expressible approximation, NOT an exact {0.5.14, 0.5.17} -# gate (CORRECTED 2026-08-18, code review): `_check_compat`'s `!=` only -# excludes the literal point version, so 0.5.15.post1/0.5.15rc1/ -# 0.5.16.post2-style variants of the excluded releases still satisfy this -# specifier (see tests/unit/collector/test_version_resolver.py's -# TestExactVersionSet for the honest, probed semantics). This is an -# accepted, narrow gap: the framework_manifest digest-pinned gate is the -# true version enforcement upstream and only ever supplies exactly 0.5.14 -# or 0.5.17 in a sanctioned run, so the leak is unreachable there. -__compat__ = "sglang>=0.5.14,<=0.5.17,!=0.5.15,!=0.5.16" +# file had none of its own before this bump). +# 0.5.21 since 2026-10-01 (H20/sm90 collector port: op_smoke + path gates in the v0.5.21 image; +# findings hopper_sglang_collector_port_0514_to_0521_2026_10_01). +# One release per file (owner policy 2026-09-20); earlier code is `git log -- ` away. +__compat__ = "sglang==0.5.21" import gc import os @@ -194,13 +185,16 @@ def run_gdn_context_benchmark( perf_filename: str, sglang_version: str, device: str = "cuda:0", + mamba_ssm_dtype: str = "float32", ): """ Benchmark GDN operations for context (prefill) phase using SGLang's Triton FLA kernels. Benchmarks: 1. causal_conv1d_fn — Conv1D over packed Q+K+V channels - 2. chunk_gated_delta_rule — GDN scan (Q, K, V, g, beta) via SGLang's vendored FLA + 2. chunk_gated_delta_rule — GDN scan (Q, K, V, g, beta) via SGLang's vendored FLA, + or flashinfer_chunk_gated_delta_rule when sglang>=0.5.21 serving defaults the + prefill backend to FlashInfer for the case (see _resolve_flashinfer_gdn_prefill) """ device = torch.device(device) torch.cuda.set_device(device) @@ -219,6 +213,7 @@ def run_gdn_context_benchmark( ) conv_weight = torch.randn(conv_channels, d_conv, dtype=dtype, device=device) + flashinfer_prefill = _resolve_flashinfer_gdn_prefill(head_k_dim, head_v_dim, mamba_ssm_dtype) successful_points = 0 failed_points = 0 @@ -244,9 +239,14 @@ def run_gdn_context_benchmark( # DSA-FUSED-KS-4G-OFFSET row. Raise instead of launching the # corrupting kernel: the async illegal access otherwise poisons # the CUDA context and aborts every remaining sweep cell. + # NOT relaxed for sglang 0.5.21: its causal_conv1d_triton casts the sequence + # offsets to tl.int64 (:95,:185), yet the >=2**31 cells still fault with + # cudaErrorIllegalAddress on H20 (0.5.21 smoke 2026-10-01, facts/upgrade_rc29/ + # sgl_smoke_r6) — silicon evidence keeps the guard; the int32 arithmetic is + # elsewhere in the kernel's per-block view strides. if total_tokens * conv_channels >= 2**31: raise ValueError( - "SGLang 0.5.14 causal_conv1d Triton kernel int32 token-offset overflow: " + "SGLang causal_conv1d Triton kernel int32 token-offset overflow: " f"total_tokens={total_tokens} * conv_channels={conv_channels} >= 2**31 " "(causal_conv1d_triton.py:373-379)" ) @@ -336,19 +336,53 @@ def run_conv1d(): ): raise RuntimeError(f"failed to persist SGLang GDN context row to {perf_filename}") - def run_gdn_scan(): - chunk_gated_delta_rule( - q, - k, - v, - g, - beta, - initial_state=recurrent_state, - initial_state_indices=state_indices, - cu_seqlens=cu_seqlens, - head_first=False, - use_qk_l2norm_in_kernel=True, - ) + if flashinfer_prefill is None: + scan_kernel_source = "chunk_gated_delta_rule" + + def run_gdn_scan(): + chunk_gated_delta_rule( + q, + k, + v, + g, + beta, + initial_state=recurrent_state, + initial_state_indices=state_indices, + cu_seqlens=cu_seqlens, + head_first=False, + use_qk_l2norm_in_kernel=True, + ) + + else: + # FlashInferGDNKernel.extend @0.5.21 (kernels/gdn_flashinfer.py:465-581): + # l2norm/layout prepare, exp(g) and beta as fp32, the state rows gathered + # (fp32 pool on SM90) and written back with index_copy_, int64 cu_seqlens, + # use_qk_l2norm_in_kernel=False because the prepare already normalized. + scan_kernel_source = "flashinfer_chunk_gated_delta_rule" + prefill_fn, qkv_prepare = flashinfer_prefill + cu_seqlens_i64 = cu_seqlens.to(torch.int64) + state_rows = state_indices.to(torch.int64) + + def run_gdn_scan(): + # gather + output buffer per call, as extend does (:523-541) + initial_state_fi = recurrent_state[state_rows].to(torch.float32) + output_state_fi = torch.empty_like(initial_state_fi) + q_fi, k_fi, v_fi = qkv_prepare(q[0], k[0], v[0]) + prefill_fn( + q=q_fi, + k=k_fi, + v=v_fi, + g=torch.exp(g[0].to(torch.float32)), + beta=beta[0].to(torch.float32), + scale=None, + initial_state=initial_state_fi, + output_final_state=True, + cu_seqlens=cu_seqlens_i64, + use_qk_l2norm_in_kernel=False, + output=None, + output_state=output_state_fi, + ) + recurrent_state.index_copy_(0, state_rows, output_state_fi.to(recurrent_state.dtype)) with benchmark_with_power( device=device, @@ -363,7 +397,7 @@ def run_gdn_scan(): version=sglang_version, device_name=torch.cuda.get_device_name(device), op_name="gdn", - kernel_source="chunk_gated_delta_rule", + kernel_source=scan_kernel_source, perf_filename=perf_filename, power_stats=results["power_stats"], ): @@ -447,6 +481,113 @@ def _resolve_flashinfer_gdn_decode(mamba_ssm_dtype: str = "float32"): return gated_delta_rule_decode_pretranspose, None +def _sglang_gdn_lanes_0521(): + """Feature-probe the sglang>=0.5.21 GDN lane surface (None on older runtimes). + + 0.5.21 moved the GDN backend to srt/layers/attention/linear/gdn_backend.py and + changed what serving runs inside the GDN role (H20 serving golden 2026-10-01, + results/pathdiff/sm90/sglang-0.5.21/gdn_{ctx,gen}_Qwen3.5-0.8B.json): + * prefill: attention_registry.py:425-430 asks flashinfer_gdn_prefill_default + (gdn_backend.py:188-234) and, when it answers "flashinfer", the extend path + runs FlashInferGDNKernel.extend (kernels/gdn_flashinfer.py:465-581): + gdn_prefill_qkv_prepare_fwd + flashinfer.gdn_prefill.chunk_gated_delta_rule + instead of the vendored FLA chunk_gated_delta_rule; + * decode: Qwen3.5 hands the backend the projected (qkvz, ba) pair + (models/qwen3_5.py:160-162, 940-953; SGLANG_ENABLE_GDN_DECODE_FUSED_PROJ_CONV + defaults True, environ.py:1800) and forward_decode runs + fused_qkvzba_causal_conv1d_update_contiguous (gdn_backend.py:668-760) in place + of the split + causal_conv1d_update. + """ + try: + from sglang.kernels.ops.attention.fla.l2norm import gdn_prefill_qkv_prepare_fwd + from sglang.kernels.ops.attention.triton_gdn_fused_proj import ( + can_use_fused_qkvzba_causal_conv1d_update_contiguous, + fused_qkvzba_causal_conv1d_update_contiguous, + ) + from sglang.srt import environ as _environ + from sglang.srt.layers.attention.linear.gdn_backend import ( # noqa: F401 — lane marker + flashinfer_gdn_prefill_default, + ) + from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import ( + _get_flashinfer_gdn_kernels, + is_flashinfer_gdn_prefill_available, + ) + except ImportError: + return None + return { + "qkv_prepare": gdn_prefill_qkv_prepare_fwd, + "fused_conv_eligible": can_use_fused_qkvzba_causal_conv1d_update_contiguous, + "fused_conv": fused_qkvzba_causal_conv1d_update_contiguous, + "fused_conv_enabled": bool(_environ.envs.SGLANG_ENABLE_GDN_DECODE_FUSED_PROJ_CONV.get()), + "flashinfer_prefill_available": is_flashinfer_gdn_prefill_available, + "flashinfer_kernels": _get_flashinfer_gdn_kernels, + } + + +def _resolve_flashinfer_gdn_prefill(head_k_dim: int, head_v_dim: int, mamba_ssm_dtype: str): + """Mirror sglang>=0.5.21's own prefill-backend default for the case. + + gdn_backend.py:188-234 flashinfer_gdn_prefill_default: FlashInfer only when + no explicit --linear-attn-prefill-backend, linear_attn_backend == triton, + no deterministic / page-major mode, SM major 9 or 10 (SM100 additionally + CUDA >= 13 and a bf16 state pool; SM90 an fp32 state pool), key/value head + dim 128, a static chunked_prefill_size within the validated chunk bound + (serving default sizing: derived, not None, dynamic chunking off), and the + FlashInfer prefill kernels importable. Everything else keeps Triton FLA, + and so does this collector. Returns ``(prefill_fn, qkv_prepare)`` or + ``None`` (FLA lane); FlashInfer missing is serving's own Triton fallback + (:230-231), so it is not an error here. + """ + lanes = _sglang_gdn_lanes_0521() + if lanes is None: + return None + sm_version = get_sm_version() + sm_major = sm_version // 10 + if sm_major not in (9, 10): + print(f" SM{sm_version}: sglang>=0.5.21 keeps Triton FLA GDN prefill (flashinfer default is SM90/SM100 only).") + return None + if sm_major == 10: + cuda_major = int(torch.version.cuda.split(".", 1)[0]) if torch.version.cuda else 0 + if cuda_major < 13 or mamba_ssm_dtype != "bfloat16": + print( + f" SM{sm_version}, cuda={torch.version.cuda}, mamba_ssm_dtype={mamba_ssm_dtype!r}: " + "sglang>=0.5.21 FlashInfer GDN prefill needs CUDA>=13 and a bf16 state pool on SM100; Triton FLA." + ) + return None + elif mamba_ssm_dtype != "float32": + print( + f" SM{sm_version}, mamba_ssm_dtype={mamba_ssm_dtype!r}: " + "FlashInfer GDN prefill needs an fp32 state pool on SM90; Triton FLA." + ) + return None + if head_k_dim != 128 or head_v_dim != 128: + print( + f" head_k_dim={head_k_dim}, head_v_dim={head_v_dim}: FlashInfer GDN prefill is 128/128 only; Triton FLA." + ) + return None + if not lanes["flashinfer_prefill_available"](): + print( + " FlashInfer GDN prefill kernels unavailable in this runtime: " + "serving falls back to Triton FLA, so does this lane." + ) + return None + _available, prefill_fn, *_rest = lanes["flashinfer_kernels"]() + return prefill_fn, lanes["qkv_prepare"] + + +def _resolve_fused_gdn_decode_conv(): + """sglang>=0.5.21 decode conv lane: the fused qkvz/ba unpack + indexed Conv1D + update (gdn_backend.py:668-760), on by default (environ.py:1800). Returns + ``(fused_fn, eligibility_fn)`` or ``None`` for the split + causal_conv1d_update + lane. When the fused kernel's own eligibility check refuses a shape, serving + logs and falls back to the unfused pair (gdn_backend.py:790-815), which the + caller mirrors and kernel_source records.""" + lanes = _sglang_gdn_lanes_0521() + if lanes is None or not lanes["fused_conv_enabled"]: + return None + return lanes["fused_conv"], lanes["fused_conv_eligible"] + + def run_gdn_generation_benchmark( d_model: int, d_conv: int, @@ -503,6 +644,7 @@ def run_gdn_generation_benchmark( # Resolve serving's exact decode dispatch: SM-major 10 + bf16 uses FI; # current FP32 recipes stay on FLA, and a required-but-missing FI raises. flashinfer_gdn_decode_fn, flashinfer_gdn_decode_error = _resolve_flashinfer_gdn_decode(mamba_ssm_dtype) + fused_decode_conv = _resolve_fused_gdn_decode_conv() successful_points = 0 failed_points = 0 failures = [] @@ -555,7 +697,36 @@ def run_gdn_generation_benchmark( "model_name": model_name, } + conv_kernel_source = "causal_conv1d_update" + fused_conv_args = None + if fused_decode_conv is not None: + # Serving input at 0.5.21: the projection outputs as laid out by Qwen3.5 + # (qwen3_5.py:940-953) — [Q|K|V|Z] of width qkv_dim + v_dim and [B|A] of + # width 2*num_v_heads, both model dtype; the kernel unpacks and runs the + # indexed Conv1D update in one launch (gdn_backend.py:749-761). + projected_qkvz = torch.randn(batch_size, conv_channels + value_dim, dtype=dtype, device=device) + projected_ba = torch.randn(batch_size, 2 * num_v_heads, dtype=dtype, device=device) + fused_conv_fn, fused_conv_eligible = fused_decode_conv + fused_conv_args = (projected_qkvz, projected_ba, conv_state, conv_weight, None, state_indices) + fused_conv_kwargs = dict( + qkv_dim=conv_channels, + v_dim=value_dim, + num_v_heads=num_v_heads, + activation="silu", + ) + eligible, reason = fused_conv_eligible(*fused_conv_args, **fused_conv_kwargs) + if eligible: + conv_kernel_source = "fused_qkvzba_causal_conv1d_update_contiguous" + else: + # gdn_backend.py:790-815: serving logs the reason and runs the unfused + # split + causal_conv1d_update for this shape. + print(f" fused GDN decode conv ineligible ({reason}); serving falls back to causal_conv1d_update") + fused_conv_args = None + def run_conv1d_update(): + if fused_conv_args is not None: + fused_conv_fn(*fused_conv_args, head_v_dim=head_v_dim, **fused_conv_kwargs) + return causal_conv1d_update( mixed_qkv, conv_state, @@ -578,7 +749,7 @@ def run_conv1d_update(): version=sglang_version, device_name=torch.cuda.get_device_name(device), op_name="gdn", - kernel_source="causal_conv1d_update", + kernel_source=conv_kernel_source, perf_filename=perf_filename, power_stats=results["power_stats"], ): @@ -820,6 +991,7 @@ def run_gdn_torch( perf_filename=perf_filename, sglang_version=sglang_version, device=device, + mamba_ssm_dtype=mamba_ssm_dtype, ) elif phase == "generation": run_gdn_generation_benchmark( diff --git a/python/aisimulate/collector/sglang/collect_gemm.py b/python/aisimulate/collector/sglang/collect_gemm.py index e35cf02e6d..f44b0af77b 100644 --- a/python/aisimulate/collector/sglang/collect_gemm.py +++ b/python/aisimulate/collector/sglang/collect_gemm.py @@ -37,31 +37,31 @@ # per layer_permissions.md); the caveat is for data consumers: fp8_block # gemm rows collected at 0.5.14 and at 0.5.17 are NOT cross-version # comparable, because they time different kernels for the identical -# (m, n, k) inputs. 0.5.15/0.5.16 stay excluded: never verified, and -# version_resolver's __compat__ grammar (AND-of-comparators only, no OR) has -# no way to express a true two-point set either -- this bounded range with -# the two untested base releases carved out via != is the closest -# expressible approximation, NOT an exact {0.5.14, 0.5.17} gate (CORRECTED -# 2026-08-18, code review): `_check_compat`'s `!=` only excludes the literal -# point version, so 0.5.15.post1/0.5.15rc1/0.5.16.post2-style variants of -# the excluded releases still satisfy this specifier (see -# tests/unit/collector/test_version_resolver.py's TestExactVersionSet for -# the honest, probed semantics). This is an accepted, narrow gap: the -# framework_manifest digest-pinned gate is the true version enforcement -# upstream and only ever supplies exactly 0.5.14 or 0.5.17 in a sanctioned -# run, so the leak is unreachable there. -# SGLang 0.5.20 (GLM-5.3-Flash model pin; audited 2026-09-30 against tag -# commit 94602c9c2b7c): serving changed WHICH kernels this module's raw -# calls map to (bf16 CuTe DSL/split-K dispatch, JIT per-token fp8 quant, -# Fp8ScaledMMOp, fp8_block Triton fallback for N%64/K%128, NVFP4 padding and -# cute-dsl activation quant), so 0.5.20 runs delegate get_gemm_test_cases/ -# run_gemm to collect_gemm_0520.py (serving-built linear layers; full audit -# there). That is a whole-implementation fork selected by the installed -# version at the bottom of this module, not a per-call shim; it lives here -# instead of a registry VersionRoute because tools/perf_database/ -# changed_ops.py fails closed on versioned registry entries. 0.5.18/0.5.19 -# are not audited and stay excluded (same != caveat as above). -__compat__ = "sglang>=0.5.14,<=0.5.20,!=0.5.15,!=0.5.16,!=0.5.18,!=0.5.19" +# (m, n, k) inputs. +# 0.5.21 since 2026-10-01 (H20/sm90 collector port: op_smoke + path gates in the v0.5.21 image; +# findings hopper_sglang_collector_port_0514_to_0521_2026_10_01). +# One release per file (owner policy 2026-09-20); earlier code is `git log -- ` away. +# The SGLang 0.5.20 fork (collect_gemm_0520.py, serving-built linear layers) +# retired with the 0.5.20 lane 2026-10-30; it is git history (main at c540ff90a (before #649)). +# FIXME(re-verify, 2026-10-30): independent serving gaps, accepted by the owner for +# the 0.5.21 tables (no re-collection). Re-check at the next bump. +# 1. bf16 GEMM on SM100/103: serving resolves the auto bf16 backend to CuTe DSL +# (quantization/unquant.py:216@v0.5.21). It first sends the 68 shapes in +# _BF16_SPLITK_TUNED_TACTICS (unquant.py:105-178: m <= 32 at GB300 TP16 and +# B300 Qwen4-Exp TP4 widths) to FlashInfer split-K or direct +# (SGLANG_ENABLE_BF16_SPLITK_GEMM defaults True, environ.py:1175; dispatch at +# unquant.py:323-347), then runs the TGV kernel only where +# use_cutedsl_bf16_gemm(m, n, k) picks it (kernels/ops/gemm/ +# cutedsl_bf16_gemm.py:1357-1387: small m, 2048 <= k <= 6144, n >= 1024); +# everywhere else serving runs F.linear like this module. +# 2. nvfp4 activation quant on SM100/103: serving quantizes with FlashInfer's +# cute-dsl backend (quantization/fp4_utils.py:24@v0.5.21); this module's +# flashinfer_fp4_quantize uses the default cuda backend. The mm_fp4 GEMM itself matches. +# 3. fp8_block on SM90 and SM100/103: serving's DeepGEMM path falls back to Triton +# (triton_w8a8_block_fp8_linear) unless weight N % 64 == 0 and K % 128 == 0 +# (quantization/fp8_utils.py:1228-1243@v0.5.21; e.g. GLM-5.3-Flash widths +# 288/3336/6416/12576/19360), while this module times DeepGEMM for every shape. +__compat__ = "sglang==0.5.21" import os import random @@ -70,6 +70,7 @@ # Set it before any SGLang imports so a task compiles only its requested M. os.environ.setdefault("SGLANG_JIT_DEEPGEMM_PRECOMPILE", "0") +# setuptools/pkg_resources is absent from recent framework images. from importlib.metadata import version as _dist_version import torch @@ -139,23 +140,6 @@ from sglang.srt.layers.quantization.fp8_utils import requant_weight_ue8m0 -def _sglang_release() -> tuple[int, ...]: - """Installed SGLang release as an int tuple (local/pre-release tags dropped). - - Collection always runs with SGLang installed (collect.py resolves the - pinned runtime from the installed version first); without the - distribution metadata (offline import tests) the legacy path is kept. - """ - from importlib.metadata import PackageNotFoundError - - from packaging.version import Version - - try: - return Version(_dist_version("sglang")).release - except PackageNotFoundError: - return (0,) - - def get_gemm_test_cases(): test_cases = [] @@ -491,12 +475,3 @@ def kernel_func(): finally: op_list.clear() torch.cuda.empty_cache() - - -# Whole-implementation version fork (see the __compat__ comment): SGLang -# 0.5.20 collects through serving-built linear layers. -if _sglang_release() >= (0, 5, 20): - from collector.sglang import collect_gemm_0520 as _fork_0520 - - get_gemm_test_cases = _fork_0520.get_gemm_test_cases - run_gemm = _fork_0520.run_gemm diff --git a/python/aisimulate/collector/sglang/collect_gemm_0520.py b/python/aisimulate/collector/sglang/collect_gemm_0520.py deleted file mode 100644 index a8001fcad7..0000000000 --- a/python/aisimulate/collector/sglang/collect_gemm_0520.py +++ /dev/null @@ -1,278 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""SGLang 0.5.20 GEMM collector (serving-built linear layers). - -At SGLang 0.5.20 the same raw kernel calls that collect_gemm.py times no -longer match serving for several dtypes, so this 0.5.20 fork (selected by collect_gemm.py) builds each -GEMM the way serving does -- a ``ReplicatedLinear`` whose quant method is -constructed from the checkpoint-style quant config -- and times -``layer.quant_method.apply`` after ``process_weights_after_loading``. The -framework therefore owns weight repacking, padding, fallbacks and kernel -selection; kernel_source records the branch the framework took. -""" - -# Audited 2026-09-30 against sglang v0.5.20 (tag commit -# 94602c9c2b7cbdb8efd5c52802dac6a1c180089e, source clone) vs v0.5.17 -# (29481685). -# Dependency pins moved sglang-kernel 0.4.5->0.4.7, flashinfer -# 0.6.15.post1->0.6.18, sgl-deep-gemm 0.1.5.post1->0.2.0, torch 2.11->2.13, -# so rows are not cross-version comparable with collect_gemm.py's. -# Serving initialization: the scheduler calls initialize_fp8_gemm_config, -# initialize_fp4_gemm_config and initialize_bf16_gemm_config -# (managers/scheduler.py:999-1001); _ensure_serving_gemm_config() does the -# same after publishing ServerArgs. -# bf16: UnquantizedLinearMethod.apply (quantization/unquant.py:460-499) -# routes bf16 through _bf16_gemm_dispatch_impl when the bf16 backend is -# cutedsl, which "auto" resolves to on SM100/103 (unquant.py:215-220). -# The dispatch order is FlashInfer split-K/direct for the tuned (m,n,k) -# table (:106-181,199-201,283-313, SGLANG_ENABLE_BF16_SPLITK_GEMM), -# CuTe DSL TGV when use_cutedsl_bf16_gemm(m,n,k) (cutedsl_bf16_gemm.py: -# 1357-1387), else F.linear (unquant.py:315-350). collect_gemm.py times -# F.linear unconditionally, which is not serving truth on SM10x. -# fp8_block: Fp8LinearMethod (quantization/fp8.py) with the auto runner -# deepgemm_w8a8_block_fp8_linear_with_fallback when JIT DeepGEMM is on -# (fp8_utils.py:799-810); shapes with N%64 or K%128 fall back to Triton -# with fp32 scales (fp8_utils.py:1094-1107) and their scales are not -# requantized to UE8M0 (fp8.py:758-770, model_loader/utils.py:265-291). -# nvfp4: ModelOptFp4LinearMethod (quantization/modelopt_quant.py:1706-2075) -# pads N/K to 32 (pad_nvfp4_weight, :195-250), quantizes activations with -# fp4_utils.fp4_quantize (FlashInfer "cute-dsl" backend on SM10x, -# fp4_utils.py:26-49) and runs mm_fp4 with the auto fp4 runner -# (fp4_utils.py:145-158: cute-dsl on SM100/103). -# The collect_gemm.py FIXME(kernel-limit) "n<128 or k<128" is false at this -# version (fp8_block falls back to Triton; nvfp4 pads to 32), so this fork -# queues those shapes and lets the framework handle them. -# Per-tensor fp8 is not collected at this version: no 0.5.20 consumer reads -# it (GLM-5.3-Flash serves fp8_block, nvfp4 and bf16 GEMMs only). -__compat__ = "sglang==0.5.20" - -import os -import random - -# SGLang reads this flag while importing ``deep_gemm_wrapper.compile_utils``. -# Set it before any SGLang imports so a task compiles only its requested M. -os.environ.setdefault("SGLANG_JIT_DEEPGEMM_PRECOMPILE", "0") - -from importlib.metadata import version as _dist_version - -import torch -from collector.case_generator import get_gemm_case_specs -from collector.helper import benchmark_with_power, get_sm_version, log_perf - -_SERVING_GEMM_CONFIG_READY = False - - -def _ensure_serving_gemm_config() -> None: - """Publish dummy ServerArgs, then run serving's GEMM backend initializers.""" - global _SERVING_GEMM_CONFIG_READY - if _SERVING_GEMM_CONFIG_READY: - return - from sglang.srt.runtime_context import get_server_args - from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler - - try: - get_server_args() # raises ValueError when nothing is published (runtime_context.py:1031-1037) - except ValueError: - set_global_server_args_for_scheduler(ServerArgs(model_path="dummy")) - from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config - from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config - from sglang.srt.layers.quantization.unquant import initialize_bf16_gemm_config - - initialize_fp8_gemm_config() - initialize_fp4_gemm_config() - initialize_bf16_gemm_config() - _SERVING_GEMM_CONFIG_READY = True - - -def get_gemm_test_cases(): - sm_version = get_sm_version() - if sm_version < 90: - gemm_list = ["bfloat16"] - elif sm_version < 100: - gemm_list = ["fp8_block", "bfloat16"] - elif sm_version < 110: - gemm_list = ["fp8_block", "bfloat16", "nvfp4"] - else: - # SM120 dense fp8_block serving is CUTLASS (fp8_utils.py:832-833); - # not collected here, unchanged from collect_gemm.py. - gemm_list = ["bfloat16", "nvfp4"] - - requested_gemm_types = os.environ.get("AIC_COLLECT_GEMM_TYPES") - if requested_gemm_types: - requested = {item.strip() for item in requested_gemm_types.split(",") if item.strip()} - gemm_list = [gemm_type for gemm_type in gemm_list if gemm_type in requested] - - test_cases = [[gemm_type, case.x, case.n, case.k] for case in get_gemm_case_specs() for gemm_type in gemm_list] - # Group DeepGEMM JIT cache hits the same way collect_gemm.py does. - random.seed(42) - random.shuffle(test_cases) - return test_cases - - -def _quant_config(gemm_type: str): - if gemm_type == "bfloat16": - return None - if gemm_type == "fp8_block": - from sglang.srt.layers.quantization.fp8 import Fp8Config - - return Fp8Config(is_checkpoint_fp8_serialized=True, activation_scheme="dynamic", weight_block_size=[128, 128]) - if gemm_type == "nvfp4": - from sglang.srt.layers.quantization.modelopt_quant import ModelOptFp4Config - - # packed_modules_mapping must be a mapping: get_quant_method's - # is_layer_skipped iterates it (quantization/utils.py:84 @v0.5.20). - return ModelOptFp4Config( - is_checkpoint_nvfp4_serialized=True, group_size=16, exclude_modules=[], packed_modules_mapping={} - ) - raise ValueError(f"no serving quant config for {gemm_type!r}") - - -def _fill_checkpoint_tensors(layer: "torch.nn.Module", gemm_type: str) -> None: - """Populate the checkpoint-native parameters create_weights allocated.""" - with torch.no_grad(): - for name, param in layer.named_parameters(recurse=False): - data = param.data - if data.dtype == torch.uint8: - data.random_(0, 256) # packed e2m1 nibbles: no NaN/Inf encodings - elif data.dtype == torch.float8_e4m3fn: - if "scale" in name: - data.copy_(torch.ones(data.shape, device=data.device).to(data.dtype)) - else: - data.copy_(((torch.rand(data.shape, device=data.device) - 0.5) * 2).to(data.dtype)) - elif data.is_floating_point(): - if "scale" in name: - data.copy_(torch.rand(data.shape, device=data.device) * 0.5 + 0.5) - else: - data.copy_(torch.randn(data.shape, device=data.device) * 0.02) - - -def _build_linear(gemm_type: str, n: int, k: int, device: str): - from sglang.srt.layers.linear import ReplicatedLinear - - layer = ReplicatedLinear( - k, n, bias=False, params_dtype=torch.bfloat16, quant_config=_quant_config(gemm_type), prefix="collector_gemm" - ).to(device) - _fill_checkpoint_tensors(layer, gemm_type) - layer.quant_method.process_weights_after_loading(layer) - layer.requires_grad_(False) - return layer - - -def _bf16_kernel_source(m: int, n: int, k: int) -> str: - """Mirror the branch unquant.UnquantizedLinearMethod.apply takes (unquant.py:460-499).""" - from sglang.srt.layers.quantization import unquant - - if not unquant.get_bf16_gemm_backend().is_cutedsl(): - return "sglang_torch_linear" - if unquant._enable_bf16_splitk_gemm and unquant.use_bf16_splitk_gemm(m, n, k): - return "sglang_flashinfer_bf16_direct" if unquant._prefer_direct(m, n, k) else "sglang_flashinfer_bf16_splitk" - if unquant._use_cutedsl_bf16_gemm is not None and unquant._use_cutedsl_bf16_gemm(m, n, k): - return "sglang_cutedsl_tgv_bf16" - return "sglang_torch_linear" - - -def _fp8_block_kernel_source(layer, n: int, k: int) -> str: - from sglang.srt.layers.quantization.fp8_utils import deepgemm_w8a8_block_fp8_linear_with_fallback - - runner = layer.quant_method.w8a8_block_fp8_linear - if runner is deepgemm_w8a8_block_fp8_linear_with_fallback: - # fp8_utils.py:1094-1107: DeepGEMM only for N%64==0 and K%128==0. - if n % 64 == 0 and k % 128 == 0: - return "sglang_deepgemm_gemm_nt_f8f8bf16" - return "sglang_triton_w8a8_block_fp8" - return f"sglang_{getattr(runner, '__name__', type(runner).__name__)}" - - -def _nvfp4_kernel_source() -> str: - from sglang.srt.layers.quantization.fp4_utils import get_fp4_gemm_runner_backend - - backend = get_fp4_gemm_runner_backend().get_flashinfer_backend() - return f"sglang_flashinfer_{str(backend).replace('-', '')}_nvfp4" - - -def run_gemm(gemm_type, batch_size, N, K, *, perf_filename, device="cuda:0"): # noqa: N803 - if gemm_type not in {"fp8_block", "bfloat16", "nvfp4"}: - raise ValueError(f"unsupported SGLang GEMM type {gemm_type!r}") - torch.cuda.set_device(device) - torch.set_default_device(device) - _ensure_serving_gemm_config() - M = batch_size # noqa: N806 - dtype = torch.bfloat16 - - def create_gemm(): - x = torch.randn((M, K), dtype=dtype, device=device) - layer = _build_linear(gemm_type, N, K, device) - - def gemm_op(): - return layer.quant_method.apply(layer, x) - - gemm_op.layer = layer - return gemm_op - - op_list = [] - try: - gpu_mem = torch.cuda.get_device_properties(device).total_memory - persistent_bytes = M * K * 2 + M * N * 2 - transient_bytes = 0 - if gemm_type == "bfloat16": - persistent_bytes += N * K * 2 - elif gemm_type == "fp8_block": - persistent_bytes += N * K + M * K + N * 4 + M * 4 - transient_bytes = N * K * 4 - elif gemm_type == "nvfp4": - persistent_bytes += int(N * K * 0.75) - transient_bytes = N * K * 2 - budget = int(gpu_mem * 0.45) - if persistent_bytes + transient_bytes >= budget: - outside_loop_count = 1 - else: - outside_loop_count = max(1, min(6, (budget - transient_bytes) // max(persistent_bytes, 1))) - - for _ in range(outside_loop_count): - op_list.append(create_gemm()) - - # Eager dry run: JIT-compiles outside graph capture. - op_list[0]() - torch.cuda.synchronize() - if gemm_type == "bfloat16": - kernel_source = _bf16_kernel_source(M, N, K) - elif gemm_type == "fp8_block": - kernel_source = _fp8_block_kernel_source(op_list[0].layer, N, K) - else: - kernel_source = _nvfp4_kernel_source() - - def kernel_func(): - for op in op_list: - op() - - nvtx_tag = f"{gemm_type}_m{M}_n{N}_k{K}" - torch.cuda.nvtx.range_push(nvtx_tag) - try: - with benchmark_with_power( - device=device, - kernel_func=kernel_func, - num_warmups=3, - num_runs=6, - repeat_n=1, - ) as results: - pass - finally: - torch.cuda.nvtx.range_pop() - - if not log_perf( - item_list=[ - {"gemm_dtype": gemm_type, "m": M, "n": N, "k": K, "latency": results["latency_ms"] / len(op_list)} - ], - framework="SGLang", - version=_dist_version("sglang"), - device_name=torch.cuda.get_device_name(device), - op_name="gemm", - kernel_source=kernel_source, - perf_filename=perf_filename, - power_stats=results["power_stats"], - ): - raise RuntimeError(f"Failed to persist SGLang GEMM performance row to {perf_filename}") - finally: - op_list.clear() - torch.cuda.empty_cache() diff --git a/python/aisimulate/collector/sglang/collect_kda.py b/python/aisimulate/collector/sglang/collect_kda.py index 4bc206f3de..c4238cc7a0 100644 --- a/python/aisimulate/collector/sglang/collect_kda.py +++ b/python/aisimulate/collector/sglang/collect_kda.py @@ -17,7 +17,8 @@ Generation (decode) phase (kda_backend.forward_decode): - causal_conv1d_update: packed single-step conv over 3P channels - - fused_recurrent_kda_packed_decode: packed KDA recurrence (T=1) + - fused_sigmoid_gating_delta_rule_update: KDA recurrence with lower_bound (T=1; + the packed kernel is skipped for lower-bound layers, kda_backend.py:737-742) - "kda_fused_decode" where serving's attempt-and-verify covered() probe accepts the shard (see _fused_decode_module below) @@ -53,21 +54,24 @@ """ # The kimi-k3 branch build (https://github.com/sgl-project/sglang/tree/kimi-k3) -# reports 0.5.16. 0.5.20 adds GLM-5.3-Flash (glm5_next) KDA, audited against -# the GB300 0.5.20 image (see the GLM section below). The file-level range -# spans both; each KDA architecture is additionally gated at runtime to its -# own audited release by _KDA_ARCHITECTURE_COMPAT (Kimi-K3 0.5.16, GLM-5.3- -# Flash 0.5.20), raising KdaRuntimeNotAuditedError otherwise. -__compat__ = "sglang>=0.5.16,<=0.5.20,!=0.5.17,!=0.5.18,!=0.5.19" +# reports 0.5.16 and was the only KDA runtime until the kernels landed in the +# official releases (0.5.21: models/kimi_linear.py, layers/attention/linear/ +# kda_backend.py with the same Triton lanes — chunk_kda, fused_recurrent_kda_ +# packed_decode / fused_sigmoid_gating_delta_rule_update, kda_fused_decode, +# causal_conv1d_fn/update — kda_backend.py:6-34, kernels/kda_triton.py:109-237; +# H20 smoke 2026-10-01, 91 cells ok, the remaining 21 are the >=2**31 conv +# guard cells). 0.5.20 added GLM-5.3-Flash (glm5_next) KDA, audited against +# the GB300 0.5.20 image (see the GLM section below). Each KDA architecture +# is additionally gated at runtime by _KDA_ARCHITECTURE_COMPAT, raising +# KdaRuntimeNotAuditedError otherwise. +# One release per file (owner policy 2026-09-20); earlier code is `git log -- ` away. +__compat__ = "sglang==0.5.21" import gc import os from typing import TYPE_CHECKING if TYPE_CHECKING: - from sglang.kernels.ops.attention.fla.fused_recurrent import ( - fused_recurrent_kda_packed_decode, - ) from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, ) @@ -213,6 +217,11 @@ def run_kda_context_benchmark( # [2**31, 3*2**31) crash with cudaErrorIllegalAddress on both # Hopper SM90 (2026-07 campaign coverage boundary) and SM100 # (B200, 2026-07-28), while every cell under this bound passes. + # NOT relaxed for sglang 0.5.21: its causal_conv1d_triton casts the sequence + # offsets to tl.int64 (:95,:185), yet the >=2**31 cells still fault with + # cudaErrorIllegalAddress on H20 (0.5.21 smoke 2026-10-01, facts/upgrade_rc29/ + # sgl_smoke_r6) — silicon evidence keeps the guard; the int32 arithmetic is + # elsewhere in the kernel's per-block view strides. if total_tokens * conv_channels >= 2**31: raise ValueError( "SGLang causal_conv1d Triton kernel int32 token-offset overflow: " @@ -406,7 +415,7 @@ def run_kda_generation_benchmark( print(f" Benchmarking batch_size={batch_size}") a = a_log = b = conv_pool = dt_bias = mixed_qkv = None - output = recurrent_state = state_indices = None + recurrent_state = state_indices = None try: num_warmups = 3 num_runs = 10 @@ -421,7 +430,6 @@ def run_kda_generation_benchmark( b = torch.randn(batch_size, num_v_heads, dtype=dtype, device=device) a_log = torch.zeros(num_v_heads, dtype=torch.float32, device=device) dt_bias = torch.ones(num_v_heads * head_k_dim, dtype=torch.float32, device=device) - output = torch.empty(batch_size, 1, num_v_heads, head_v_dim, dtype=dtype, device=device) common_log_data = { "phase": "generation", @@ -543,24 +551,43 @@ def run_conv1d_update(): ): raise RuntimeError(f"failed to persist SGLang KDA generation row to {perf_filename}") - def run_kda_packed_decode(): - fused_recurrent_kda_packed_decode( - mixed_qkv=mixed_qkv, - a=a, - b=b, + # kda_backend.py:737-742 @0.5.21 skips the packed T=1 kernel for a layer with a + # lower_bound (Kimi-K3: KDA_LOWER_BOUND); :771-786: after the conv update, q/k/v are the + # [Q|K|V] split viewed as 1 x n x h x d, and the plain Triton decode runs + # fused_sigmoid_gating_delta_rule_update over cu_seqlens with lower_bound. + recurrence_kernel_source = "fused_sigmoid_gating_delta_rule_update" + plain_q, plain_k, plain_v = ( + t.unflatten(-1, (-1, dim)).unsqueeze(0) + for t, dim in zip( + mixed_qkv.split([proj_size] * 3, dim=-1), (head_k_dim, head_k_dim, head_v_dim), strict=True + ) + ) + plain_a = a.view(1, batch_size, num_v_heads, head_k_dim) + plain_b = b.view(1, batch_size, num_v_heads) + plain_cu_seqlens = torch.arange(0, batch_size + 1, dtype=torch.int32, device=device) + + def run_kda_decode(): + fused_sigmoid_gating_delta_rule_update( A_log=a_log, dt_bias=dt_bias, - scale=head_k_dim**-0.5, - initial_state=recurrent_state, - out=output, - ssm_state_indices=state_indices, + q=plain_q, + k=plain_k, + v=plain_v, + a=plain_a, + b=plain_b, + initial_state_source=recurrent_state, + initial_state_indices=state_indices, + cu_seqlens=plain_cu_seqlens, use_qk_l2norm_in_kernel=True, + softplus_beta=1.0, + softplus_threshold=20.0, + is_kda=True, lower_bound=KDA_LOWER_BOUND, ) with benchmark_with_power( device=device, - kernel_func=run_kda_packed_decode, + kernel_func=run_kda_decode, num_warmups=num_warmups, num_runs=num_runs, repeat_n=1, @@ -571,7 +598,7 @@ def run_kda_packed_decode(): version=sglang_version, device_name=torch.cuda.get_device_name(device), op_name="kda", - kernel_source="fused_recurrent_kda_packed_decode", + kernel_source=recurrence_kernel_source, perf_filename=perf_filename, power_stats=results["power_stats"], ): @@ -585,7 +612,7 @@ def run_kda_packed_decode(): continue finally: a = a_log = b = conv_pool = dt_bias = mixed_qkv = None - output = recurrent_state = state_indices = None + recurrent_state = state_indices = None _cleanup("generation") summary = f"ok={successful_points} error={failed_points} skip=0" @@ -951,8 +978,11 @@ def _is_glm5_next_kda(model_name: str) -> bool: # replicates; any other installed release raises KdaRuntimeNotAuditedError (a # classified failure) instead of timing a possibly different kernel path. _KDA_ARCHITECTURE_COMPAT = { - KIMI_K3_KDA_ARCHITECTURE: "sglang==0.5.16", - GLM5_NEXT_KDA_ARCHITECTURE: "sglang==0.5.20", + KIMI_K3_KDA_ARCHITECTURE: "sglang==0.5.21", + # GLM-5.3-Flash (0.5.20 only) retired 2026-10-30: its route is git history + # (main at c540ff90a (before #649)). On 0.5.21 serving fuses the unquantized + # qkvbfg/bfg projections (strided views), so the GLM path needs a port + # before it is re-admitted; GLM rows raise KdaRuntimeNotAuditedError. } @@ -968,7 +998,9 @@ def _kda_architecture(model_name: str) -> str: def _require_audited_runtime(model_name: str, runtime_version: str) -> None: architecture = _kda_architecture(model_name) - compat = _KDA_ARCHITECTURE_COMPAT[architecture] + compat = _KDA_ARCHITECTURE_COMPAT.get(architecture) + if compat is None: + raise KdaRuntimeNotAuditedError(f"sglang {runtime_version} is not an audited KDA runtime for {architecture}") if not _check_compat(compat, runtime_version): raise KdaRuntimeNotAuditedError( f"sglang {runtime_version} is not an audited KDA runtime for {architecture} (audited: {compat})" @@ -1325,9 +1357,6 @@ def run_kda_torch( contextlib.redirect_stdout(_devnull_file), contextlib.redirect_stderr(_devnull_file), ): - from sglang.kernels.ops.attention.fla.fused_recurrent import ( - fused_recurrent_kda_packed_decode, - ) from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, ) @@ -1341,7 +1370,6 @@ def run_kda_torch( "causal_conv1d_fn": causal_conv1d_fn, "causal_conv1d_update": causal_conv1d_update, "chunk_kda": chunk_kda, - "fused_recurrent_kda_packed_decode": fused_recurrent_kda_packed_decode, "fused_sigmoid_gating_delta_rule_update": fused_sigmoid_gating_delta_rule_update, } ) diff --git a/python/aisimulate/collector/sglang/collect_mhc_module.py b/python/aisimulate/collector/sglang/collect_mhc_module.py index 97be868c57..e308ebeeb3 100644 --- a/python/aisimulate/collector/sglang/collect_mhc_module.py +++ b/python/aisimulate/collector/sglang/collect_mhc_module.py @@ -8,18 +8,18 @@ # capable image or put a matching SGLang source tree on PYTHONPATH. from __future__ import annotations -# The file-level range spans the two audited routes; each architecture is -# additionally gated to its own audited release by _ARCHITECTURE_COMPAT, so -# the DeepSeek-V4 path still runs only on 0.5.14 (its ModelRunner/ServerArgs -# API use is 0.5.14-specific) and GLM only on 0.5.20. -__compat__ = "sglang>=0.5.14,<=0.5.20" +# 0.5.21 added 2026-10-01 (H20/sm90 collector port: op_smoke + path gates in the v0.5.21 image; +# findings hopper_sglang_collector_port_0514_to_0521_2026_10_01). +# Each architecture is additionally gated by _ARCHITECTURE_COMPAT +# (DeepSeek-V4 on 0.5.21 via the native DeepseekV4Config). +# One release per file (owner policy 2026-09-20); earlier code is `git log -- ` away. +__compat__ = "sglang==0.5.21" import argparse import copy import gc import json import os -import random import sys import tempfile from collections.abc import Sequence @@ -57,8 +57,11 @@ # Per-architecture audited SGLang releases; anything else raises # MhcRuntimeNotAuditedError (a classified failure) instead of running. _ARCHITECTURE_COMPAT = { - DSV4_ARCHITECTURE: "sglang==0.5.14", - GLM5NEXT_ARCHITECTURE: "sglang==0.5.20", + DSV4_ARCHITECTURE: "sglang==0.5.21", + # GLM-5.3-Flash (0.5.20 only) retired 2026-10-30: its route is git history + # (main at c540ff90a (before #649)). 0.5.21 moved MHCState to layer_boundary, removed + # ParallelState/init_torch_distributed and fuses post+pre for T<=16, so the + # GLM path needs a port before it is re-admitted. } # GLM-5.3-Flash mHC call sites in SGLang 0.5.20 serving @@ -224,7 +227,13 @@ def _patched_model_dir(model_id: str) -> str: # Match collect_dsv4_attn.py: current Transformers does not know a # native deepseek_v4 config, while SGLang selects the V4 model class from # the architectures field. - config["model_type"] = "deepseek_v3" + try: + # sglang>=0.5.21: native DeepseekV4Config (V4-only fields such as hc_pre_from_prev_sublayer) + import sglang.srt.configs.deepseek_v4 # noqa: F401 + + config["model_type"] = "deepseek_v4" + except ImportError: + config["model_type"] = "deepseek_v3" tmp_dir = os.path.join( tempfile.gettempdir(), @@ -283,18 +292,18 @@ def _load_one_layer_runner( _set_envs_and_config(server_args) model_config = ModelConfig.from_server_args(server_args) + from collector.sglang.runtime_compat import free_local_port, init_runtime_config + + nccl_port = free_local_port() + + _runner_parallel_kwargs = init_runtime_config(server_args, gpu_id, nccl_port=nccl_port, model_config=model_config) return ModelRunner( model_config=model_config, mem_fraction_static=mem_fraction_static, gpu_id=gpu_id, - tp_rank=0, - tp_size=1, - pp_rank=0, - pp_size=1, - moe_ep_rank=0, - moe_ep_size=1, - nccl_port=29500 + random.randint(0, 10000), + nccl_port=nccl_port, server_args=server_args, + **_runner_parallel_kwargs, ) @@ -547,11 +556,11 @@ def run_mhc_module( gc.collect() finally: del model_runner + from collector.sglang.runtime_compat import reset_expert_location_metadata from sglang.srt.distributed.parallel_state import ( destroy_distributed_environment, destroy_model_parallel, ) - from sglang.srt.eplb import expert_location as _expert_location # Mirror SGLang 0.5.14 cleanup_dist_env_and_memory: destroying only the # torch process group leaves parallel_state._WORLD pointing at a dead @@ -562,11 +571,11 @@ def run_mhc_module( # and fail the worker rather than hiding retained groups. destroy_model_parallel() destroy_distributed_environment() - # SGLang has no public reset for this module global (its serving - # process never re-creates a ModelRunner); set_global_... asserts None, - # so a second in-worker task fails init unless it is returned to the - # module's pre-init state here. - _expert_location._global_expert_location_metadata = None + # SGLang has no public reset for the expert-location metadata (its + # serving process never re-creates a ModelRunner); set_global_... asserts + # None, so a second in-worker task fails init unless it is returned to + # its pre-init state here. + reset_expert_location_metadata() torch.cuda.empty_cache() gc.collect() summary = f"ok={len(results)} error={error_count} skip=0 total={len(results) + error_count}" @@ -587,6 +596,7 @@ def _init_glm5next_runtime(model_path: str, *, device: str, mem_fraction_static: (model_runner.py L456-458). No model weights are loaded: the mHC kernels only need this process state plus the layer's own parameters. """ + from collector.sglang.runtime_compat import free_local_port from sglang.srt.configs.model_config import ModelConfig from sglang.srt.distributed import bootstrap from sglang.srt.distributed.parallel_state_wrapper import ParallelState @@ -629,7 +639,7 @@ def _init_glm5next_runtime(model_path: str, *, device: str, mem_fraction_static: model_config=model_config, device="cuda", ps=ParallelState.trivial(gpu_id=gpu_id), - dist_port=29500 + random.randint(0, 10000), + dist_port=free_local_port(), is_draft_worker=False, local_omp_cpuid=None, ) diff --git a/python/aisimulate/collector/sglang/collect_mla.py b/python/aisimulate/collector/sglang/collect_mla.py index 9391b32e59..9415391dfd 100644 --- a/python/aisimulate/collector/sglang/collect_mla.py +++ b/python/aisimulate/collector/sglang/collect_mla.py @@ -9,13 +9,18 @@ mocking, runtime dispatch, and perf logging. """ -__compat__ = "sglang==0.5.14" +# 0.5.21 since 2026-10-01 (H20/sm90 collector port: op_smoke + path gates in the v0.5.21 image; +# findings hopper_sglang_collector_port_0514_to_0521_2026_10_01). +# One release per file (owner policy 2026-09-20); earlier code is `git log -- ` away. +__compat__ = "sglang==0.5.21" +import contextlib +import inspect import math import os import random +from importlib.metadata import version as _dist_version # setuptools/pkg_resources is absent from sglang>=0.5.21 images -import pkg_resources import sglang.srt.layers.dp_attention import sglang.srt.server_args import torch @@ -93,6 +98,9 @@ def _mla_compute_dtype(backend: str, kv_cache_dtype: torch.dtype) -> str: class MockModelConfig: + def get_max_num_attention_heads(self) -> int: # model_config.py:1504 @0.5.21 (triton_backend.py:255) + return int(self.num_attention_heads) + def __init__( self, context_len: int = 32768, @@ -120,10 +128,26 @@ def __init__( self.head_dim = 256 self.v_head_dim = v_head_dim self.hf_text_config = self + # sglang>=0.5.21 TritonAttnBackend.__init__ (the sm89 / sm120 MLA default) asks + # mambaish_config(model_config) whether this is a hybrid linear-attention arch + # (triton_backend.py:287 -> configs/hybrid_arch.py reads hf_config / get_text_config() + # / linear_attn_registry_result); an MLA module is none of them. fa3 / trtllm_mla + # (sm90 / sm100) never reach that probe, so only the Triton lane needs this. + self.linear_attn_registry_result = None + self.is_draft_model = False + + class _MockHFConfig: + architectures = ["DeepseekV3ForCausalLM"] # noqa: RUF012 + model_type = "deepseek_v3" + + def get_text_config(self): + return self + + self.hf_config = _MockHFConfig() self.scaling = scaling self.is_local_attention_model = False - def get_num_kv_heads(self, tp_size: int): + def get_num_kv_heads(self, tp_size: int, dcp_size: int = 1): # model_config.py:1512 @0.5.21 return 1 @@ -160,9 +184,14 @@ def __init__(self, kv_cache_dtype: torch.dtype, page_size: int): # stays off exactly as in production (triton_backend.py:214-217). # The SM90 fa3 / SM100 trtllm_mla init paths never read this field. self.cuda_graph_config = None + self.enable_prefill_cp = False # sglang 0.5.16 flashattention_backend.py:51 + self.enable_dp_attention = False # same site: prefill-CP/DP-attention gate class MockModelRunner: + def decode_num_tokens_per_req(self, *, num_draft_tokens=None) -> int: + return 1 + def __init__( self, device: torch.device, @@ -175,6 +204,8 @@ def __init__( self.gpu_id = device.index if device.index is not None else torch.cuda.current_device() self.tp_size = 1 self.kv_cache_dtype = kv_cache_dtype + # sglang>=0.5.21 backends read the runner's server-arg spelling (model_runner.py:1481) + self.kv_cache_dtype_str = "fp8_e4m3" if kv_cache_dtype == torch.float8_e4m3fn else "auto" self.dtype = torch.bfloat16 self.page_size = page_size self.req_to_token_pool = None @@ -190,7 +221,48 @@ def __init__( # Keep attributes for compatibility across sglang versions (older code ignores them) self.is_hybrid_swa = self.model_config.is_hybrid_swa self.attn_cp_size = 1 # Context parallelism size; required by FlashAttentionBackend in sglang >=0.5.10 + # sglang 0.5.21 TRTLLMMLABackend (the MLA backend sglang picks on SM100, _get_backends) reads + # model_runner.max_running_requests (workspace sizing) and model_runner.is_draft_worker; + # on Hopper fa3/flashinfer never asked, so the attributes were missing until the first + # Blackwell smoke (B200 2026-10-04: AttributeError 'MockModelRunner' has no attribute + # 'max_running_requests'). 512 = the generator's --max-running-requests default. + self.max_running_requests = 512 + self.is_draft_worker = False + # sglang 0.5.16 reads the parallel geometry from model_runner.ps + # (flashattention_backend.py:183,271-274; parallel_state_wrapper.py:6-24) + try: + from sglang.srt.distributed.parallel_state_wrapper import ParallelState + + self.ps = ParallelState.trivial(gpu_id=0) + except ImportError: + from types import SimpleNamespace + + self.ps = SimpleNamespace( + tp_rank=0, + tp_size=1, + pp_rank=0, + pp_size=1, + dp_rank=0, + dp_size=1, + attn_tp_rank=0, + attn_tp_size=1, + attn_cp_rank=0, + attn_cp_size=1, + attn_dp_rank=0, + attn_dp_size=1, + moe_ep_rank=0, + moe_ep_size=1, + moe_dp_rank=0, + moe_dp_size=1, + dcp_size=1, + gpu_id=0, + ) self.server_args = MockServerArgs(kv_cache_dtype, page_size) + self.is_draft_worker = False + # sglang>=0.5.21 backends ask the runner for logits rows per decode slot + # (model_runner.py:796 decode_num_tokens_per_req; triton_backend.py:218,222); + # the kernel collectors never run speculative decoding -> 1 + self.spec_algorithm = None # read by the 0.5.16 attention backends (draft-worker branches) self.use_mla_backend = True @@ -214,8 +286,91 @@ def create_req_to_token_pool( return pool, token_matrix.contiguous() +def _extend_seq_lens_cpu(forward_batch): + """CPU int32 mirror of the per-request extend lengths of a plain EXTEND batch, else None. + + The values are sglang's own ``forward_batch.extend_seq_lens_cpu`` — the list its + trtllm_mla backend already reads for ``max_seq_len`` (trtllm_mla_backend.py:850 + @v0.5.21) — so nothing is reconstructed, only mirrored into a CPU tensor. + """ + if forward_batch.forward_mode != ForwardMode.EXTEND: + return None + extend_seq_lens = getattr(forward_batch, "extend_seq_lens_cpu", None) + if not extend_seq_lens: + return None + return torch.tensor([int(n) for n in extend_seq_lens], dtype=torch.int32, device="cpu") + + +@contextlib.contextmanager +def _ragged_prefill_with_cpu_seq_lens(seq_lens_cpu): + """Hand flashinfer the CPU length mirrors that keep trtllm_ragged_attention_deepseek capturable. + + flashinfer 0.6.18 (sglang 0.5.21's pin) takes optional ``q_seq_lens_cpu`` / + ``kv_seq_lens_cpu`` (prefill.py:4797-4798). Without them the wrapper derives + the per-row lengths from the device indptrs with a host sync and refuses to + run under CUDA-graph capture ("q_seq_lens_cpu and kv_seq_lens_cpu must be + provided during CUDA graph capture", prefill.py:5083-5091); with them it + validates the mirrors on the host (prefill.py:5053-5077) and launches the + same trtllm-gen kernel with the same device arguments. Older flashinfer + builds without the parameters never synchronized here and are left alone. + """ + import flashinfer.prefill as flashinfer_prefill + + original = flashinfer_prefill.trtllm_ragged_attention_deepseek + parameters = inspect.signature(original).parameters + if "q_seq_lens_cpu" not in parameters or "kv_seq_lens_cpu" not in parameters: + yield + return + + def _with_cpu_seq_lens(*args, **kwargs): + kwargs.setdefault("q_seq_lens_cpu", seq_lens_cpu) + kwargs.setdefault("kv_seq_lens_cpu", seq_lens_cpu) + return original(*args, **kwargs) + + flashinfer_prefill.trtllm_ragged_attention_deepseek = _with_cpu_seq_lens + try: + yield + finally: + flashinfer_prefill.trtllm_ragged_attention_deepseek = original + + +class _GraphCapturableTRTLLMMLABackend(TRTLLMMLABackend): + """sglang's trtllm_mla backend with a CUDA-graph-capturable prefill kernel call. + + Serving never graph-captures prefill, so sglang 0.5.21's + ``TRTLLMMLABackend._run_prefill_kernel`` (trtllm_mla_backend.py:1098-1141) + calls ``flashinfer.prefill.trtllm_ragged_attention_deepseek`` without the + CPU length mirrors and flashinfer refuses the capture (see + ``_ragged_prefill_with_cpu_seq_lens``). The collector times the kernel under + graph replay — eager timing of this lane floors at the Python launch loop + (0.13 ms on B200, 0.29 ms on GB200 for every isl <= 1024 cell; data + validation 2026-10-11) — so the no-prefix prefill call gets the mirrors + here. This changes how the kernel is invoked, never which kernel runs: the + chunked-prefix calls (their own ``cum_seq_lens_kv``) and every other + backend path go through sglang unchanged. + """ + + _prefill_seq_lens_cpu = None + + def init_forward_metadata(self, forward_batch): + super().init_forward_metadata(forward_batch) + self._prefill_seq_lens_cpu = _extend_seq_lens_cpu(forward_batch) + + def _run_prefill_kernel(self, *args, **kwargs): + seq_lens_cpu = self._prefill_seq_lens_cpu + cum_seq_lens_q = kwargs.get("cum_seq_lens_q") + cum_seq_lens_kv = kwargs.get("cum_seq_lens_kv") + if seq_lens_cpu is None or cum_seq_lens_q is None or cum_seq_lens_kv is not cum_seq_lens_q: + return super()._run_prefill_kernel(*args, **kwargs) + with _ragged_prefill_with_cpu_seq_lens(seq_lens_cpu): + return super()._run_prefill_kernel(*args, **kwargs) + + def benchmark_layer(layer, forward_batch, q, k, v, q_rope, k_rope, **kwargs): - # Use benchmark_with_power context manager + # CUDA-graph replay for every lane: a capture failure raises (helper default + # allow_graph_fail=False) and an eager result is refused below, so a kernel + # that stops being capturable fails the case instead of shipping its launch + # floor (sglang mla_context on B200/GB200, data validation 2026-10-11). device = q.device def kernel_func(): @@ -236,6 +391,8 @@ def kernel_func(): ) as results: pass + if results.get("used_cuda_graph") is not True: + raise RuntimeError("SGLang MLA requires CUDA graph timing; refusing to publish eager timing") return results["latency_ms"], results["power_stats"] @@ -349,6 +506,9 @@ def run_mla( perf_filename, device="cuda:0", ): + from collector.sglang.runtime_compat import ensure_offline_runtime_published + + ensure_offline_runtime_published() # sglang>=0.5.20 backends read get_exec()/get_parallel() torch.cuda.set_device(device) torch_device = torch.device(device) random.seed(0) @@ -363,6 +523,28 @@ def run_mla( expected_page_size = 64 if selected_backend == "trtllm_mla" else 1 if tokens_per_block != expected_page_size: raise ValueError(f"SGLang {selected_backend} requires page_size={expected_page_size}, got {tokens_per_block}") + # FIXME(kernel-limit): sglang 0.5.21 trtllm_mla DECODE on Blackwell with 64 < local heads < 128. + # flashinfer's trtllm-gen MLA decode refuses that head range (flashinfer mla/_core.py:3948) and + # falls back to the cute-dsl runner, but sglang 0.5.21's trtllm_mla_backend passes + # multi_ctas_kv_counter_buffer unconditionally on Blackwell -> "multi_ctas_kv_counter_buffer is + # only supported when a trtllm-gen runner is selected" (_core.py:4278). Serving hits the same + # raise (B200 identity probe: Kimi-K3 bf16 tp1, 96 heads; finding + # sm100_probe_failures_recheck_2026_10_04 (i)), so the cell is serving-unreachable at this tp; + # the b200_sxm shard run (GitLab job 469988017) failed exactly the 362 num_heads=96 decode cases + # (bf16 + fp8 KV) and nothing else. Prefill (mla_context, 96 heads) is unaffected. Fail-closed: + # recorded as a classified failure, never a predicted skip or a silent backend swap. + if ( + not is_context_phase + and selected_backend == "trtllm_mla" + and get_sm_version() in {100, 103} + and 64 < local_num_heads < 128 + ): + raise RuntimeError( + "FIXME(kernel-limit): sglang 0.5.21 trtllm_mla decode on SM100/103 does not support " + f"64 < local_num_heads < 128 (num_heads={num_heads}, tp_size={tp_size} -> {local_num_heads}); " + "flashinfer trtllm-gen MLA decode rejects the head range and the cute-dsl fallback rejects " + "multi_ctas_kv_counter_buffer (mla/_core.py:3948, :4278); serving fails identically (Kimi-K3 tp1)" + ) model_runner = MockModelRunner( torch_device, @@ -385,7 +567,12 @@ def run_mla( model_runner.server_args.prefill_attention_backend = selected_backend model_runner.server_args.decode_attention_backend = selected_backend # Set global args after potential overrides. - sglang.srt.server_args.set_global_server_args_for_scheduler(model_runner.server_args) + try: + from sglang.srt.runtime_context import publish as _publish # noqa: F401 + except ImportError: + sglang.srt.server_args.set_global_server_args_for_scheduler(model_runner.server_args) + # >=0.5.20: set_global_server_args_for_scheduler publishes (resolve_once on a real ServerArgs record); + # ensure_offline_runtime_published() already installed the process record, the mock drives the backends directly. # Define dimensions based on phase kv_lora_rank = KV_LORA_RANK @@ -432,6 +619,9 @@ def run_mla( enable_memory_saver=False, ) model_runner.token_to_kv_pool = kv_pool + from collector.sglang.runtime_compat import attach_kv_index_translator + + attach_kv_index_translator(model_runner) if selected_backend == "trtllm_mla": # TRTLLMMLABackend inherits FlashInferMLAAttnBackend which creates @@ -440,7 +630,7 @@ def run_mla( import gc gc.collect() - attn_backend = TRTLLMMLABackend(model_runner) + attn_backend = _GraphCapturableTRTLLMMLABackend(model_runner) kernel_source = "trtllm_mla" elif selected_backend == "triton": attn_backend = TritonAttnBackend(model_runner) @@ -628,7 +818,7 @@ def run_mla( } ], framework="SGLang", - version=pkg_resources.get_distribution("sglang").version, + version=_dist_version("sglang"), device_name=torch.cuda.get_device_name(device), op_name=f"mla_{'context' if is_context_phase else 'generation'}", kernel_source=kernel_source, diff --git a/python/aisimulate/collector/sglang/collect_mla_bmm.py b/python/aisimulate/collector/sglang/collect_mla_bmm.py index 3a253d4791..b34bd5f0c0 100644 --- a/python/aisimulate/collector/sglang/collect_mla_bmm.py +++ b/python/aisimulate/collector/sglang/collect_mla_bmm.py @@ -8,16 +8,32 @@ shapes, selects SGLang kernel helpers, and logs the resulting MLA BMM perf rows. """ -__compat__ = "sglang==0.5.14" +# 0.5.21 since 2026-10-01 (H20/sm90 collector port: op_smoke + path gates in the v0.5.21 image; +# findings hopper_sglang_collector_port_0514_to_0521_2026_10_01). +# One release per file (owner policy 2026-09-20); earlier code is `git log -- ` away. +__compat__ = "sglang==0.5.21" + +from importlib.metadata import version as _dist_version # setuptools/pkg_resources is absent from sglang>=0.5.21 images -import pkg_resources import torch from collector.case_generator import get_mla_bmm_case_specs from collector.helper import benchmark_with_power, get_sm_version, log_perf -from sgl_kernel import bmm_fp8 -from sglang.srt.layers.quantization.fp8_kernel import ( - per_tensor_quant_mla_fp8, -) + +try: + from sgl_kernel import bmm_fp8 + + _BMM_FP8_SOURCE = "sglang_sgl_kernel_bmm_fp8" +except ImportError: # sglang>=0.5.21: sgl_kernel no longer exports bmm_fp8; serving wraps flashinfer's + # (layers/quantization/fp8_utils.py:246-273 @0.5.21, custom op flashinfer_bmm_fp8_batched) + from sglang.srt.layers.quantization.fp8_utils import bmm_fp8 + + _BMM_FP8_SOURCE = "sglang_flashinfer_bmm_fp8" +try: + from sglang.srt.layers.quantization.fp8_kernel import per_tensor_quant_mla_fp8 +except ModuleNotFoundError: # sglang>=0.5.21: the Triton quant kernels moved to sglang.kernels.ops + # (kernels/ops/quantization/fp8_kernel.py:956 @0.5.21; same kernel, the serving MLA fp8 lane + # imports it from there — deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py:8) + from sglang.kernels.ops.quantization.fp8_kernel import per_tensor_quant_mla_fp8 def _supported_dtypes() -> set[str]: @@ -97,10 +113,10 @@ def kernel_func(): } ], framework="SGLang", - version=pkg_resources.get_distribution("sglang").version, + version=_dist_version("sglang"), device_name=torch.cuda.get_device_name(device), op_name="mla_gen_pre", - kernel_source="sglang_sgl_kernel_bmm_fp8" if dtype == "fp8" else "sglang_torch_bmm", + kernel_source=_BMM_FP8_SOURCE if dtype == "fp8" else "sglang_torch_bmm", perf_filename=perf_filename, power_stats=results["power_stats"], ): @@ -181,10 +197,10 @@ def kernel_func(): } ], framework="SGLang", - version=pkg_resources.get_distribution("sglang").version, + version=_dist_version("sglang"), device_name=torch.cuda.get_device_name(device), op_name="mla_gen_post", - kernel_source="sglang_sgl_kernel_bmm_fp8" if dtype == "fp8" else "sglang_torch_bmm", + kernel_source=_BMM_FP8_SOURCE if dtype == "fp8" else "sglang_torch_bmm", perf_filename=perf_filename, power_stats=results["power_stats"], ): diff --git a/python/aisimulate/collector/sglang/collect_mla_module.py b/python/aisimulate/collector/sglang/collect_mla_module.py index b1132a3708..96b8f7430a 100644 --- a/python/aisimulate/collector/sglang/collect_mla_module.py +++ b/python/aisimulate/collector/sglang/collect_mla_module.py @@ -22,7 +22,10 @@ python collect_mla_module.py --mode generation --attn-type mla """ -__compat__ = "sglang==0.5.14" +# 0.5.21 since 2026-10-01 (H20/sm90 collector port: op_smoke + path gates in the v0.5.21 image; +# findings hopper_sglang_collector_port_0514_to_0521_2026_10_01). +# One release per file (owner policy 2026-09-20); earlier code is `git log -- ` away. +__compat__ = "sglang==0.5.21" import argparse import gc @@ -44,12 +47,27 @@ from importlib.metadata import version as get_version import torch +from collector.sglang.runtime_compat import resolved_arg # 0.5.21: derived server args live behind the resolving view try: - from helper import benchmark_with_power, get_sm_version, log_perf + from helper import ( + PhaseClock, + benchmark_with_power, + get_sm_version, + log_perf, + render_phase_totals, + run_measurement_subprocess, + ) except ModuleNotFoundError: sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - from helper import benchmark_with_power, get_sm_version, log_perf + from helper import ( + PhaseClock, + benchmark_with_power, + get_sm_version, + log_perf, + render_phase_totals, + run_measurement_subprocess, + ) try: from collector.case_generator import ( @@ -130,6 +148,10 @@ _NATIVE_NVFP4_SMS = {100, 103, 120} +# Per-subprocess sum of the point phase times printed by run_mla_module. +_POINT_PHASE_TOTALS: dict[str, float] = {} + + def _env_flag(name: str, default: bool = False) -> bool: value = os.environ.get(name) if value is None: @@ -186,8 +208,8 @@ def _is_glm5_dsa_model(model_id: str) -> bool: # Set by run_mla_module() at the start of each benchmark subprocess from its -# skip_indexer arg (threaded down from run_mla_module_worker, which derives it -# from the op's perf_filename). Process-local: each subprocess runs exactly one +# skip_indexer arg (threaded down from run_mla_module_worker, which receives it +# from the registry entry's run_kwargs). Process-local: each subprocess runs exactly one # run_mla_module, so this is never shared across cases. Replaces the old # AIC_DSA_SKIP_INDEXER env that previously crossed the subprocess boundary. _SKIP_INDEXER_PASS = False @@ -218,7 +240,7 @@ def _generation_cuda_graph_enabled_for_tokens(model_runner, num_tokens: int) -> that coverage using sglang's own settings -- no AIC env override -- so the decode benchmark uses graph timing exactly where serve would. """ - decode_config = model_runner.server_args.cuda_graph_config.decode + decode_config = resolved_arg(model_runner.server_args, "cuda_graph_config").decode if decode_config.backend == "disabled": return False capture_bs = decode_config.bs @@ -236,19 +258,21 @@ def _initialize_dsa_history(model_runner, attention_module, forward_batch, histo including partial pages, rather than exposing zeros or an earlier case's contents. These inputs represent no particular checkpoint distribution. - Native contracts at SGLang 0.5.14, commit 49e384ce9d304648e9959666ecb8ce8cd98d0deb: - dsa/dsa_indexer.py:1259-1329 owns index-K quantization and opaque layout; - dsa_backend.py:1564-1577 writes latent/rotary KV using the attention layer; - memory_pool.py:2235-2286 owns BF16/FP8 MLA layouts and their scale fields. - Paths above are relative to python/sglang/srt/layers/attention except - memory_pool.py, which is in python/sglang/srt/mem_cache. + Native contracts at SGLang 0.5.21 (ported from #405's 0.5.14 @49e384ce): + dsa/dsa_indexer.py:1498-1578 owns index-K quantization and opaque layout + (its pool comes from ForwardContext, :1519); dsa_backend.py:1997-2010 writes + latent/rotary KV using the attention layer; memory_pool.py:4570-4642 owns + BF16/FP8 MLA layouts and their scale fields. Paths above are relative to + python/sglang/srt/layers/attention except memory_pool.py, which is in + python/sglang/srt/mem_cache. """ if history_length < 0: raise ValueError("DSA historical length must be nonnegative") if not history_length: return - from sglang.srt.layers.attention.dsa.triton_kernel import act_quant + # 0.5.21 moved the DSA Triton kernels out of srt (kernels/ops/attention/dsa/triton_kernel.py:86) + from sglang.kernels.ops.attention.dsa.triton_kernel import act_quant forward_context_type, forward_context = _import_sglang_forward_context() pool = model_runner.token_to_kv_pool @@ -291,21 +315,24 @@ def random_activations(*shape): def _dsa_forward_input_relay(model_runner, batch): - """Own the native scheduler relay needed by SGLang 0.5.14 input staging. - - At 49e384ce, managers/overlap_utils.py:65-97 materializes input IDs; - scheduler.py:3269-3280 uses it and FutureMap.stash even without overlap. - Attention-only execution skips embeddings, but ForwardBatch still needs - these IDs to derive the actual token count and native graph eligibility. + """Own the native scheduler relay needed by SGLang input staging. + + At 0.5.21 prepare_for_extend leaves input_ids None and stages + prefill_input_ids_cpu (schedule_batch.py:2943-2944); resolve_forward_inputs + materializes them (managers/overlap_utils.py:87-120), and the scheduler runs + it and FutureMap.stash even without overlap (scheduler.py:4588-4597, + 4674-4696). Attention-only execution skips embeddings, but ForwardBatch still + needs these IDs to derive the actual token count (forward_batch_info.py:1064-1071). + The relay is built as the scheduler builds it (scheduler.py:1622-1637). """ - from sglang.srt.managers.overlap_utils import FutureMap, decide_needs_cpu_seq_lens + from sglang.srt.managers.overlap_utils import decide_needs_confidence_relay, decide_needs_cpu_seq_lens if not hasattr(model_runner, "_aic_dsa_input_relay"): - model_runner._aic_dsa_input_relay = FutureMap( - device=batch.device, - spec_algo=batch.spec_algorithm, - req_to_token_pool=model_runner.req_to_token_pool, - needs_cpu_seq_lens=decide_needs_cpu_seq_lens(model_runner.server_args, [model_runner.attn_backend]), + model_runner._aic_dsa_input_relay = batch.spec_algorithm.create_future_map( + batch.device, + model_runner.req_to_token_pool, + needs_cpu_seq_lens=decide_needs_cpu_seq_lens((model_runner.attn_backend,)), + needs_confidence_relay=decide_needs_confidence_relay(), ) return model_runner._aic_dsa_input_relay @@ -317,9 +344,11 @@ def _validate_dsa_forward_tokens(forward_batch, expected_tokens: int) -> None: if value is None or value.numel() != expected_tokens: actual = None if value is None else value.numel() raise ValueError(f"DSA {name} has {actual} tokens; expected {expected_tokens}") - if forward_batch.num_token_non_padded_cpu != expected_tokens: + # 0.5.21 renamed the host count (forward_batch_info.py:673, set at :1071) + if forward_batch.global_num_token_non_padded_cpu != expected_tokens: raise ValueError( - f"DSA num_token_non_padded_cpu={forward_batch.num_token_non_padded_cpu}; expected {expected_tokens}" + f"DSA global_num_token_non_padded_cpu={forward_batch.global_num_token_non_padded_cpu}; " + f"expected {expected_tokens}" ) @@ -704,8 +733,8 @@ def _dsa_context_prefix_shape_is_valid( ) -> bool: """Return whether a DSA prefix-context sample is structurally valid. - The scheduler's forward mode distinguishes prefill from decode, including - a one-token extension. Native graph dispatch handles its padded bucket. + A one-token extension stays out at 0.5.21: the collector times DSA prefill + eagerly, where it illegal-accessed (runtime_limits.dsa_indexer_prefill_shape_is_supported). **Per-request context ceiling (SGLang mechanism, not an empirical constant).** The DSA/NSA indexer builds its rotary cos/sin cache for @@ -973,7 +1002,7 @@ def get_dsa_context_module_skip_indexer_test_cases(): patched out (GLM-5.2 index_topk_freq>1 reuse layers). Same shapes as the full context module — the skip behaviour is applied in - the subprocess (run_func detects the skip_indexer perf_filename). Only emit + the subprocess (the registry binds skip_indexer=True via run_kwargs). Only emit cases for models that actually share the index across layers (index_topk_freq > 1); for freq==1 models the skip layer == full layer, so a separate file would just duplicate dsa_context_module. @@ -1020,10 +1049,9 @@ def cleanup_distributed(): if hasattr(parallel_state, var_name): setattr(parallel_state, var_name, None) - import sglang.srt.eplb.expert_location as expert_location + from collector.sglang.runtime_compat import reset_expert_location_metadata - if hasattr(expert_location, "_global_expert_location_metadata"): - expert_location._global_expert_location_metadata = None + reset_expert_location_metadata() def _ensure_fp8_block_quant_config(hf_cfg) -> None: @@ -1089,11 +1117,14 @@ def _patch_nsa_rope_contiguity(model_runner): def _make_contiguous_forward(orig): def _forward(positions, query, key, fused_set_kv_buffer_arg=None): + # by keyword: the 4th positional of forward_cuda is offsets + # (rotary_embedding/base.py:380-387 @0.5.21) + kwargs = {} if fused_set_kv_buffer_arg is None else {"fused_set_kv_buffer_arg": fused_set_kv_buffer_arg} return orig( positions.contiguous(), query.contiguous(), key.contiguous() if key is not None else key, - fused_set_kv_buffer_arg, + **kwargs, ) return _forward @@ -1248,13 +1279,8 @@ def load_model_runner( "Marlin is valid only for INT4-WO" ) - import random - from sglang.srt.configs.model_config import ModelConfig from sglang.srt.entrypoints.engine import _set_envs_and_config - from sglang.srt.layers.moe import initialize_moe_config - from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config - from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.server_args import ServerArgs from sglang.srt.utils import suppress_other_loggers @@ -1289,7 +1315,11 @@ def load_model_runner( tp_size=1, trust_remote_code=True, disable_radix_cache=True, - disable_prefill_cuda_graph=attention_backend != "dsa", + # FIXME(re-verify): 0.5.21 serves DSA prefill under the "breakable" prefill graph + # (cuda_graph_config.py:122-131,160-162; buckets 4..2048, cuda_graph_hook.py:622-639), + # not the 0.5.14 tc_piecewise runner dsa_prefill_graph.py adapts, so the collector keeps + # prefill graphs off and times DSA prefill eagerly, as the 0.5.21 tables were measured. + disable_prefill_cuda_graph=True, kv_cache_dtype=sglang_kv_dtype, max_total_tokens=max_total_tokens, chunked_prefill_size=chunked_prefill_size, @@ -1329,11 +1359,10 @@ def load_model_runner( server_args.json_model_override_args = json.dumps(override_args) _set_envs_and_config(server_args) - initialize_moe_config(server_args) - initialize_fp8_gemm_config(server_args) - initialize_fp4_gemm_config(server_args) - nccl_port = 29500 + random.randint(0, 10000) + gpu_id * 100 + from collector.sglang.runtime_compat import free_local_port + + nccl_port = free_local_port() model_config = ModelConfig.from_server_args(server_args) expected_architecture = _module_model_architecture(model_path) @@ -1353,18 +1382,19 @@ def load_model_runner( if native_quant == "fp8_block": _ensure_fp8_block_quant_config(model_config.hf_config) + # parallel geometry / config publish / distributed bootstrap: one helper per sglang generation + from collector.sglang.runtime_compat import init_runtime_config + + _runner_parallel_kwargs = init_runtime_config( + server_args, gpu_id, nccl_port=nccl_port, model_config=model_config, tp_rank=gpu_id, tp_size=server_args.tp_size + ) model_runner = ModelRunner( model_config=model_config, mem_fraction_static=server_args.mem_fraction_static, gpu_id=gpu_id, - tp_rank=gpu_id, - tp_size=server_args.tp_size, - pp_rank=0, - pp_size=1, - moe_ep_rank=0, - moe_ep_size=1, nccl_port=nccl_port, server_args=server_args, + **_runner_parallel_kwargs, ) model_runner.alloc_memory_pool() @@ -1466,9 +1496,8 @@ def dummy_qkv_latent_func(h, fb): log_mla_dtype = "fp8" if backend_name == "trtllm_mla" and kv_cache_dtype == "fp8" else "bfloat16" if ordinary_mla or attn_type == "dsa": - # SGLang 0.5.14 @ 49e384ce9d304648e9959666ecb8ce8cd98d0deb: - # deepseek_v2.py:1902-1921,2068 supplies the native projection to - # AttentionInputs for both MLA and DSA. Keep its fused/quantized + # SGLang 0.5.21 deepseek_v2.py:2462-2487,2622-2626 supplies the + # native projection to AttentionInputs for both MLA and DSA. Keep its fused/quantized # dispatch; the module perf row includes this down-projection. dummy_qkv_latent_func = attention_module.prepare_qkv_latent @@ -1589,7 +1618,10 @@ def _run_prefill( from array import array is_wideep_mla = attn_type == "mla" and not ordinary_mla - from sglang.srt.layers.communicator import AttentionInputs, get_attn_tp_context + try: # sglang>=0.5.21: layers/layer_boundary/adapters/attention.py:45,167 + from sglang.srt.layers.layer_boundary.adapters.attention import AttentionInputs, get_attn_tp_context + except ImportError: + from sglang.srt.layers.communicator import AttentionInputs, get_attn_tp_context from sglang.srt.managers.schedule_batch import Req, ScheduleBatch from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.chunk_cache import ChunkCache @@ -1601,6 +1633,7 @@ def _run_prefill( forward_context_type, forward_context = _import_sglang_forward_context() print(f"\nPrefill: batch_size={batch_size}, seq_length={seq_length}, prefix_len={prefix_len}") + point_clock = PhaseClock() native_graph_context = None try: @@ -1621,9 +1654,18 @@ def _run_prefill( req.prefix_indices = prefix_indices[i] req.full_untruncated_fill_ids = array("q", req.origin_input_ids) req.fill_len = full_length - req.set_extend_input_len(seq_length if prefix_len else full_length) + # sglang 0.5.16 replaced Req.set_extend_input_len(n) with + # set_extend_range(start, end) (schedule_batch.py:1153: the extend + # window is the last n of fill_len tokens); keep the old call on + # pins that still have it. + _n = seq_length if prefix_len else full_length + if hasattr(req, "set_extend_input_len"): + req.set_extend_input_len(_n) + else: + req.set_extend_range(full_length - _n, full_length) req.logprob_start_len = 0 reqs.append(req) + point_clock.mark("reqs") cache_params = CacheInitParams( disable=True, @@ -1644,16 +1686,29 @@ def _run_prefill( ) with _temporarily_chunked_alloc_extend(model_runner, batch_size * seq_length): batch.prepare_for_extend() + point_clock.mark("batch") if attn_type == "dsa": from sglang.srt.managers.overlap_utils import resolve_forward_inputs resolve_forward_inputs(batch, _dsa_forward_input_relay(model_runner, batch)) - forward_batch = ForwardBatch.init_new(batch, model_runner) + # sglang 0.5.16 made return_hidden_states_before_norm a required + # keyword (ForwardBatch.init_new); older pins do not accept it — + # same dispatch collect_msa_module.py:1015 uses. + import inspect as _inspect + + _fb_kw = ( + {"return_hidden_states_before_norm": False} + if "return_hidden_states_before_norm" in _inspect.signature(ForwardBatch.init_new).parameters + else {} + ) + forward_batch = ForwardBatch.init_new(batch, model_runner, **_fb_kw) if attn_type == "dsa": _validate_dsa_forward_tokens(forward_batch, batch_size * seq_length) model_runner.attn_backend.init_forward_metadata(forward_batch) + point_clock.mark("metadata") if attn_type == "dsa": _initialize_dsa_history(model_runner, attention_module, forward_batch, prefix_len) + point_clock.mark("history") hidden_states = torch.randn( batch_size * seq_length, @@ -2053,8 +2108,9 @@ def _skip_kwargs(): def call_attention_module(): if ordinary_mla or attn_type == "dsa": - # communicator.py:235-246 caches QKV within AttentionInputs. - # Serving constructs fresh inputs per layer/forward (:685-689). + # AttentionInputs caches QKV (layer_boundary/adapters/attention.py:69-80 + # @0.5.21); serving builds fresh inputs per layer/forward + # (layer_boundary/prepare.py:179-189). # Reset inside EVERY timed iteration, including warmup, so the # down-projection is measured rather than reused from warmup. get_attn_tp_context().set_attn_inputs( @@ -2094,7 +2150,7 @@ def call_attention_module(): # is built in init_forward_metadata, not via the forward context). # Metadata was planned once for this batch above. Serving # does this before model.forward, not once per layer: - # 49e384ce, model_executor/runner/eager_runner.py:246-262. + # model_executor/runner/eager_runner.py:293-321 then :370-376 @0.5.21. # Retain the native ForwardContext while timing the layer. return attention_module( positions=positions, @@ -2112,6 +2168,7 @@ def call_attention_module(): **_skip_kwargs(), ) + point_clock.mark("setup") last_can_run_graph = None def call_full_model_piecewise_replay(): @@ -2140,7 +2197,8 @@ def call_full_model_piecewise_replay(): call_target = graph_context.__enter__() native_graph_context = graph_context # Native piecewise metadata disables MHA_ONE_SHOT even for - # small inputs (dsa_backend.py:2436-2444 @ 49e384ce). + # small inputs (dsa_backend.py:2436-2444 @ 49e384ce; at 0.5.21 + # tc_piecewise or breakable, dsa_backend.py:3603-3606). # Label the path selected after native load_batch, not the # earlier eager preparation or the requested backend. if model_runner.attn_backend.use_mha: @@ -2201,6 +2259,7 @@ def call_full_model_piecewise_replay(): "refusing to record a skip row with full-indexer latency." ) + point_clock.mark("warmup") module_cuda_graph = None if use_module_cuda_graph: try: @@ -2214,6 +2273,7 @@ def call_full_model_piecewise_replay(): print(" Module CUDA graph capture failed") raise + point_clock.mark("capture") # Timed runs — run ALL iterations back-to-back, ONE sync at the end, then # divide by count. Per-iteration cuda.synchronize would insert GPU-idle # bubbles + per-launch latency between modules (NOT representative); the @@ -2231,6 +2291,7 @@ def call_full_model_piecewise_replay(): end_event.record() torch.cuda.synchronize() avg_time_ms = start_event.elapsed_time(end_event) / num_iterations + point_clock.mark("timed") # Log perf — wideep MLA uses old filename/op_name/kernel_source conventions try: @@ -2318,6 +2379,11 @@ def call_full_model_piecewise_replay(): cleanup_errors.append(f"{cleanup_name}: {type(cleanup_exc).__name__}: {cleanup_exc}") if cleanup_errors: raise RuntimeError(f"DSA/MLA prefill cleanup failed: {'; '.join(cleanup_errors)}") + point_clock.mark("cleanup") + point_clock.add_to(_POINT_PHASE_TOTALS) + print( + f" [DSA point timing] prefill b={batch_size}, s={seq_length}, prefix={prefix_len}: {point_clock.render()}" + ) def _run_decode( @@ -2348,7 +2414,10 @@ def _run_decode( from array import array is_wideep_mla = attn_type == "mla" - from sglang.srt.layers.communicator import AttentionInputs, get_attn_tp_context + try: # sglang>=0.5.21: layers/layer_boundary/adapters/attention.py:45,167 + from sglang.srt.layers.layer_boundary.adapters.attention import AttentionInputs, get_attn_tp_context + except ImportError: + from sglang.srt.layers.communicator import AttentionInputs, get_attn_tp_context from sglang.srt.managers.schedule_batch import Req, ScheduleBatch from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.chunk_cache import ChunkCache @@ -2361,6 +2430,7 @@ def _run_decode( forward_context_type, forward_context = _import_sglang_forward_context() print(f"\nDecode: batch_size={batch_size}, kv_cache_length={seq_length}") + point_clock = PhaseClock() try: model_runner.req_to_token_pool.clear() @@ -2377,11 +2447,15 @@ def _run_decode( req.prefix_indices = torch.empty((0,), dtype=torch.int64) req.full_untruncated_fill_ids = array("q", req.origin_input_ids) req.fill_len = len(req.origin_input_ids) - req.set_extend_input_len(req.fill_len) + if hasattr(req, "set_extend_input_len"): + req.set_extend_input_len(req.fill_len) + else: # sglang 0.5.16: schedule_batch.py:1153 set_extend_range(start, end) + req.set_extend_range(0, req.fill_len) req.logprob_start_len = 0 req.cached_tokens = 0 req.already_computed = 0 reqs.append(req) + point_clock.mark("reqs") cache_params = CacheInitParams( disable=True, @@ -2410,19 +2484,35 @@ def _run_decode( for req in batch.reqs: req.output_ids.append(0) if attn_type == "dsa": - # Match the non-overlap scheduler's sampled-token handoff. The - # synthetic next token is the same zero appended to each Req. - input_relay.stash(batch.req_pool_indices, torch.zeros(batch_size, dtype=torch.int64, device=batch.device)) + # Match the non-overlap scheduler's sampled-token handoff + # (scheduler.py:4588-4597, 4688-4696 @0.5.21). The synthetic next + # token is the same zero appended to each Req. + from sglang.srt.managers.overlap_utils import RelayPayload + + input_relay.stash( + batch.req_pool_indices, + RelayPayload(bonus_tokens=torch.zeros(batch_size, dtype=torch.int64, device=batch.device)), + ) batch.input_ids = None batch.prepare_for_decode() + point_clock.mark("batch") if attn_type == "dsa": resolve_forward_inputs(batch, input_relay) - forward_batch_decode = ForwardBatch.init_new(batch, model_runner) + import inspect as _inspect + + _fb_kw = ( + {"return_hidden_states_before_norm": False} + if "return_hidden_states_before_norm" in _inspect.signature(ForwardBatch.init_new).parameters + else {} + ) # sglang 0.5.16 required keyword (see the context path above) + forward_batch_decode = ForwardBatch.init_new(batch, model_runner, **_fb_kw) if attn_type == "dsa": _validate_dsa_forward_tokens(forward_batch_decode, batch_size) model_runner.attn_backend.init_forward_metadata(forward_batch_decode) + point_clock.mark("metadata") if attn_type == "dsa": _initialize_dsa_history(model_runner, attention_module, forward_batch_decode, seq_length) + point_clock.mark("history") decode_hidden = torch.randn( batch_size, @@ -2435,6 +2525,7 @@ def _run_decode( attn_inputs_decode = AttentionInputs(decode_hidden, forward_batch_decode, dummy_qkv_latent_func) get_attn_tp_context().set_attn_inputs(attn_inputs_decode) + point_clock.mark("inputs") if use_module_cuda_graph or _env_flag("AIC_ENABLE_MODULE_CUDA_GRAPH"): _patch_nsa_indexer_compile_for_module_cuda_graph(attention_module) @@ -2477,8 +2568,8 @@ def kernel_func(): capture_context = model_capture_mode() if use_benchmark_cuda_graph else nullcontext() with capture_context, forward_context(forward_context_type(attn_backend=model_runner.attn_backend)): if attn_type == "dsa": - # SGLang 0.5.14 @ 49e384ce, communicator.py:235-246, - # 685-689: QKV is cached within one AttentionInputs only. + # SGLang 0.5.21 layer_boundary/adapters/attention.py:69-80, + # prepare.py:179-189: QKV is cached within one AttentionInputs only. # Refresh for every forward so capture includes the native # projection instead of reusing the warmup's latent. get_attn_tp_context().set_attn_inputs( @@ -2521,6 +2612,7 @@ def kernel_func(): if _skip_indexer: attention_module.next_skip_topk = original_next_skip_topk + point_clock.mark("prewarm") print(f" Decode module CUDA graph: {use_benchmark_cuda_graph} (tokens={batch_size})") with benchmark_with_power( @@ -2532,6 +2624,8 @@ def kernel_func(): use_cuda_graph=use_benchmark_cuda_graph, ) as results: pass + point_clock.mark("bench") + point_clock.absorb(results.get("phase_seconds"), prefix="bench_") if _skip_indexer and _skip_state["prev_topk"] is None: raise RuntimeError( @@ -2627,6 +2721,9 @@ def kernel_func(): cleanup_errors.append(f"{cleanup_name}: {type(cleanup_exc).__name__}: {cleanup_exc}") if cleanup_errors: raise RuntimeError(f"DSA/MLA decode cleanup failed: {'; '.join(cleanup_errors)}") + point_clock.mark("cleanup") + point_clock.add_to(_POINT_PHASE_TOTALS) + print(f" [DSA point timing] decode b={batch_size}, kv={seq_length}: {point_clock.render()}") def _resolve_perf_path(output_path: str | None, filename: str) -> str: @@ -2666,9 +2763,9 @@ def run_mla_module( (batch_size, seq_length) combos for the specified phase. ``skip_indexer`` selects the GLM-5.2 reuse-layer variant (the per-layer - indexer is patched out and rows are tagged ``_skip_indexer``). It mirrors - ``is_prefill``: the worker derives it from the op's perf_filename and passes - it down. Recorded in the ``_SKIP_INDEXER_PASS`` process global so the + indexer is patched out and rows are tagged ``_skip_indexer``). The worker + receives it from the registry entry's run_kwargs and passes it down. + Recorded in the ``_SKIP_INDEXER_PASS`` process global so the existing ``_dsa_skip_indexer_enabled`` call sites need no signature change. """ global _SKIP_INDEXER_PASS @@ -2924,6 +3021,7 @@ def run_mla_module( error_count = len(cases) - logged_count summary = f"ok={logged_count} error={error_count} skip=0 total={len(cases)}" print(f"[{attn_type.upper()}] {phase_name.lower()} {summary}") + print(f"[{attn_type.upper()} timing] {phase_name.lower()}: {render_phase_totals(_POINT_PHASE_TOTALS)}") if error_count > 0: raise RuntimeError( f"{attn_type.upper()} module {phase_name.lower()} failed strict completeness: {summary}; " @@ -2972,30 +3070,29 @@ def _run_mla_subprocess( f"chunked_prefill_size={chunked_prefill_size!r})\n" ) - proc = subprocess.Popen( - [sys.executable, "-c", code], - env=env, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - cwd=os.path.dirname(os.path.abspath(__file__)), - ) - # Per-subprocess timeout is OFF by default. A DSA context subprocess sweeps # the full isl x prefix grid for one batch_size and can legitimately run for # many minutes; a wall-clock cap silently truncated the sweep (killed mid-grid, # kept the partial rows, logged a WARNING not an error) and left the perf DB - # with short prefix coverage, forcing AIC to extrapolate. Unset = - # communicate(timeout=None) = block until the subprocess finishes on its own. + # with short prefix coverage, forcing AIC to extrapolate. Unset = block until + # the subprocess finishes on its own; the executor's stall watchdog still sees + # liveness, because the wait beats while the child prints its points, and a + # reclaimed worker takes the child down with it (helper.run_measurement_subprocess). # Users can still opt into a cap by setting AIC_MLA_MODULE_SUBPROCESS_TIMEOUT_SEC. _timeout_env = os.environ.get("AIC_MLA_MODULE_SUBPROCESS_TIMEOUT_SEC") subprocess_timeout = int(_timeout_env) if _timeout_env else None try: - stdout, _ = proc.communicate(timeout=subprocess_timeout) + returncode, stdout = run_measurement_subprocess( + [sys.executable, "-c", code], + env=env, + cwd=os.path.dirname(os.path.abspath(__file__)), + timeout=subprocess_timeout, + ) if stdout: print(stdout.decode("utf-8", errors="replace")) except subprocess.TimeoutExpired as exc: - proc.kill() - proc.wait() + if exc.output: + print(exc.output.decode("utf-8", errors="replace")) raise RuntimeError( f"{attn_type.upper()} module {phase} subprocess timed out " f"after {subprocess_timeout}s " @@ -3003,7 +3100,7 @@ def _run_mla_subprocess( f"gemm={gemm_type}, target_tp={target_tp_size})" ) from exc - if proc.returncode != 0: + if returncode != 0: # Include last lines of subprocess output in the error so the cause # is visible in the error group (subprocess stderr is merged into stdout). tail = "" @@ -3011,7 +3108,7 @@ def _run_mla_subprocess( lines = stdout.decode("utf-8", errors="replace").strip().splitlines() tail = "\n".join(lines[-30:]) # last 30 lines should include traceback raise RuntimeError( - f"{attn_type.upper()} module {phase} subprocess failed (exit code {proc.returncode})\n" + f"{attn_type.upper()} module {phase} subprocess failed (exit code {returncode})\n" f"--- subprocess output (last 30 lines) ---\n{tail}" ) @@ -3032,6 +3129,7 @@ def run_mla_module_worker( perf_filename: str, device: str = "cuda:0", chunked_prefill_size: int | None = None, + skip_indexer: bool = False, ): """Worker-compatible wrapper used by collector/collect.py. @@ -3045,7 +3143,10 @@ def run_mla_module_worker( For DSA test cases, it defaults to None and _get_backends() is used. perf_filename and device are keyword-only arguments supplied by - collect.py via functools.partial and the worker dispatch loop. + collect.py via functools.partial and the worker dispatch loop; + skip_indexer is bound the same way from the registry entry's run_kwargs + (dsa_*_module_skip_indexer ops), so the full and skip ops share one + table and one run_func and differ only in this explicit argument. """ device_str = str(device) if not isinstance(device, str) else device gpu_id = int(device_str.split(":")[-1]) if ":" in device_str else 0 @@ -3056,9 +3157,9 @@ def run_mla_module_worker( # module; the only difference is (1) the per-layer indexer (mqa+topk) is # patched out via skip_topk in _run_prefill/_run_generation and (2) the rows # are tagged with an op_name "_skip_indexer" suffix in the same perf file. - # Derived from the op's perf_filename and threaded down as an explicit arg — - # exactly like is_prefill above; no env var crosses the subprocess boundary. - skip_indexer = "skip_indexer" in os.path.basename(perf_filename) + # Threaded down to the subprocess as an explicit arg; no env var crosses + # the subprocess boundary and nothing is inferred from the table name. + skip_indexer = bool(skip_indexer) print(f"\n{'=' * 60}") print( diff --git a/python/aisimulate/collector/sglang/collect_moe.py b/python/aisimulate/collector/sglang/collect_moe.py index d8d76999b2..441b280abf 100644 --- a/python/aisimulate/collector/sglang/collect_moe.py +++ b/python/aisimulate/collector/sglang/collect_moe.py @@ -42,19 +42,10 @@ # _pin_moe_runner_backend below for the full trace and the version-branched # fix: 0.5.14 keeps the exact original global read/write/restore, 0.5.17 # goes through the new RuntimeContext/Flags singleton via its own sanctioned -# override() primitive). 0.5.15/0.5.16 stay excluded: never verified, and -# version_resolver's __compat__ grammar (AND-of-comparators only, no OR) has -# no way to express a true two-point set either -- this bounded range with -# the two untested base releases carved out via != is the closest -# expressible approximation, NOT an exact {0.5.14, 0.5.17} gate (CORRECTED -# 2026-08-18, code review): `_check_compat`'s `!=` only excludes the literal -# point version, so 0.5.15.post1/0.5.15rc1/0.5.16.post2-style variants of -# the excluded releases still satisfy this specifier (see -# tests/unit/collector/test_version_resolver.py's TestExactVersionSet for -# the honest, probed semantics). This is an accepted, narrow gap: the -# framework_manifest digest-pinned gate is the true version enforcement -# upstream and only ever supplies exactly 0.5.14 or 0.5.17 in a sanctioned -# run, so the leak is unreachable there. +# override() primitive). +# 0.5.21 since 2026-10-01 (H20/sm90 collector port: op_smoke + path gates in the v0.5.21 image; +# findings hopper_sglang_collector_port_0514_to_0521_2026_10_01). +# One release per file (owner policy 2026-09-20); earlier code is `git log -- ` away. # 0.5.20 audit (2026-09-30; source tag v0.5.20 == 94602c9c2b7c vs v0.5.17, # source-only). 0.5.18/0.5.19 # were never verified and are carved out with != (same caveat as above). @@ -88,14 +79,14 @@ # (moe_runner/flashinfer_cutlass.py:17,208-209), which the single-rank # parallel patch below does not replace, so flashinfer_cutlass lanes (the # SM120 nvfp4 map) are expected to fail until that module is added. -__compat__ = "sglang>=0.5.14,<=0.5.20,!=0.5.15,!=0.5.16,!=0.5.18,!=0.5.19" +__compat__ = "sglang==0.5.21" import gc import importlib -import itertools import os import sys import tempfile +from collections.abc import Callable from contextlib import contextmanager # The SGLang 0.5.20 image ships setuptools without pkg_resources; the @@ -282,6 +273,7 @@ balanced_logits, benchmark_with_power, build_rank0_local_workload, + flashinfer_autotuned, get_sm_version, log_perf, power_law_logits_v3, @@ -304,6 +296,7 @@ balanced_logits, benchmark_with_power, build_rank0_local_workload, + flashinfer_autotuned, get_sm_version, log_perf, power_law_logits_v3, @@ -345,6 +338,13 @@ def _ensure_writable_flashinfer_cubin_dir() -> None: def get_moe_test_cases(): + """One case per (quant mode, shape, tp, ep, distribution) carrying the whole token + sweep, like the vLLM and TRT-LLM MoE collectors: the framework layer (weights, quant + processing, backend pins) is built once per case and every token point is timed and + persisted inside it. The previous one-case-per-token layout rebuilt the layer for + each of the 27 token points, which made the SGLang nvfp4 tables ~2x the fixed cost of + the vLLM ones (2026-10-28 re-collection: 4.4-5.8 h per SGLang shard vs 1-1.4 h per + vLLM shard on the same shapes).""" sm_version = get_sm_version() moe_list = get_moe_quantization_modes("sglang", sm_version=sm_version) @@ -358,7 +358,7 @@ def get_moe_test_cases(): model_name = common_moe_testcase.model_name num_tokens_list = common_moe_testcase.num_tokens_list - for moe_type, num_tokens in itertools.product(moe_list, num_tokens_list): + for moe_type in moe_list: if not moe_model_allows_quantization("sglang", model_name, moe_type): quant_policy_drops[model_name] = quant_policy_drops.get(model_name, 0) + 1 continue @@ -370,7 +370,7 @@ def get_moe_test_cases(): moe_backend = get_sglang_moe_backend(common_moe_testcase, moe_type, sm_version) base_case = [ moe_type, - num_tokens, + list(num_tokens_list), common_moe_testcase.hidden_size, common_moe_testcase.inter_size, common_moe_testcase.topk, @@ -399,7 +399,7 @@ def get_moe_test_cases(): ] physical_key = ( moe_type, - num_tokens, + tuple(num_tokens_list), common_moe_testcase.hidden_size, common_moe_testcase.inter_size, common_moe_testcase.topk, @@ -636,6 +636,7 @@ def _patch_framework_moe_parallel(*, moe_tp_size: int, moe_ep_size: int): parallel = SimpleNamespace( tp_size=moe_tp_size, tp_rank=0, + tp_group=None, # 0.5.21 topk.py:720 use_symmetric_memory(get_parallel().tp_group, disabled=True) only reads it moe_tp_size=moe_tp_size, moe_tp_rank=0, moe_ep_size=moe_ep_size, @@ -653,6 +654,11 @@ def replace(module, name, value): "sglang.srt.layers.moe.moe_runner.flashinfer_trtllm", "sglang.srt.layers.moe.moe_runner.flashinfer_mxfp4", "sglang.srt.layers.quantization.mxfp4_flashinfer_trtllm_moe", + # 0.5.21 flashinfer_cutlass runner reads get_parallel().tp_group / + # is_allocation_symmetric() at module level (moe_runner/flashinfer_cutlass.py:21, + # :27, :283, :510) — every Nemotron-3 BF16 flashinfer_cutlass case raised + # "parallel name 'tp_group' has not been written" in the 2026-10-02 H20 campaign. + "sglang.srt.layers.moe.moe_runner.flashinfer_cutlass", ): try: modules.append(importlib.import_module(module_name)) @@ -662,15 +668,9 @@ def replace(module, name, value): for module in modules: replace(module, "get_tp_group", lambda: None) replace(module, "is_allocation_symmetric", lambda: False) - if module in { - _moe_layer_mod, - _std_dispatch_mod, - _topk_mod, - _mxint4_mod, - _mxfp4_mod, - _fp8_mod, - _modelopt_mod, - }: + # Any module that imported get_parallel at module level reads our + # single-process stand-in (the dynamically imported runners included). + if hasattr(module, "get_parallel"): replace(module, "get_parallel", lambda: parallel) replace(_moe_layer_mod, "get_moe_expert_parallel_world_size", lambda: moe_ep_size) replace(_moe_layer_mod, "get_moe_expert_parallel_rank", lambda: 0) @@ -813,7 +813,8 @@ def _benchmark_framework_quantized_moe( *, moe_type: str, moe_backend: str, - num_tokens: int, + num_tokens_list: list[int], + on_result: Callable[[int, float, dict, str], None], hidden_size: int, inter_size: int, topk: int, @@ -840,8 +841,14 @@ def _benchmark_framework_quantized_moe( is_fp4_experts: bool, int4_group_size: int, device: str, -) -> tuple[float, dict, str]: - """Benchmark model-aware paths through SGLang FusedMoE (0.5.14/0.5.17).""" +) -> str: + """Benchmark model-aware paths through SGLang FusedMoE (0.5.14/0.5.17). + + The layer is constructed and its weights processed once; each token point in + ``num_tokens_list`` is then timed on it and handed to ``on_result(num_tokens, + latency_ms, power_stats, kernel_source)`` as soon as it is measured, so a crash + later in the sweep keeps the rows already persisted. Returns the kernel source. + """ if moe_backend.startswith("flashinfer"): _ensure_writable_flashinfer_cubin_dir() @@ -1011,11 +1018,21 @@ def _benchmark_framework_quantized_moe( f"requested={moe_backend}, actual={actual_backend}" ) if actual_backend == "flashinfer_mxfp4": - if quant_method._fi_kernel != "cutlass_sm90": + # Mxfp4MoEMethod dispatches flashinfer_mxfp4 to one of three leaves by GPU + # (mxfp4.py: SM90 cutlass_fused_moe w4-group-scaling, SM100 trtllm_fp4_block_scale_moe, + # SM120 cutlass MXFP8xMXFP4); name the kernel source after the leaf that actually ran + # (B200 2026-10-05: the Hopper-only check rejected the 'trtllm_sm100' leaf gpt-oss + # serves with on Blackwell — identity records: Mxfp4MoE -> trtllm_gen_moe). + _fi_leaf_sources = { + "cutlass_sm90": "sglang_flashinfer_cutlass_moe", + "trtllm_sm100": "sglang_mxfp4_flashinfer_trtllm_moe", + "flashinfer_cutlass_sm120": "sglang_flashinfer_cutlass_moe_sm120", + } + if quant_method._fi_kernel not in _fi_leaf_sources: raise RuntimeError( f"SGLang {moe_type} has unexpected FlashInfer MXFP4 leaf {quant_method._fi_kernel!r}" ) - kernel_source = "sglang_flashinfer_cutlass_moe" + kernel_source = _fi_leaf_sources[quant_method._fi_kernel] else: kernel_source = source_by_backend[actual_backend] elif ( @@ -1071,55 +1088,82 @@ def _benchmark_framework_quantized_moe( correction_bias = ( torch.zeros(num_experts, dtype=torch.float32, device=device) if has_correction_bias else None ) - topk_layer = TopK( - top_k=topk, - layer_id=0, - use_grouped_topk=num_expert_group is not None and topk_group is not None, - num_expert_group=num_expert_group, - topk_group=topk_group, - renormalize=renormalize, - scoring_func=scoring_func, - correction_bias=correction_bias, - output_format=( - TopKOutputFormat.STANDARD - if routing_method_type is None and moe_backend not in {"flashinfer_mxfp4", "triton_kernel"} - else None - ), - routed_scaling_factor=routed_scaling_factor, - apply_routed_scaling_factor_on_output=( - routed_scaling_factor is not None and moe_layer.should_fuse_routed_scaling_factor_in_topk - ), - is_fp4_experts=is_fp4_experts, - ) - hidden_states = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=device) - if distributed == "balanced": - logits = [ - balanced_logits(num_tokens, num_experts, topk).to(device=device, dtype=torch.float32) - for _ in range(5) - ] - elif distributed == "power_law": - logits = [ - power_law_logits_v3(num_tokens, num_experts, topk, moe_ep_size, power_law_alpha).to( - device=device, dtype=torch.float32 - ) - for _ in range(5) - ] + if routing_method_type == "Llama4": + # Serving builds the Llama-4 router with a custom routing function + # (fast_topk on the raw logits, then sigmoid; renormalize=False) — + # models/llama4.py:78-92 (function) and :114-118 (TopK construction) + # @0.5.21, unchanged since 0.5.14. The generic sigmoid TopK the + # collector used before routes through moe_fused_gate on 0.5.21, which + # asserts "bias is required for non-softmax routing" (kernels/ops/moe/ + # moe_fused_gate.py:384) — a kernel serving never runs for Llama-4. + from sglang.srt.models.llama4 import Llama4MoE + + topk_layer = TopK( + top_k=topk, + layer_id=0, + renormalize=False, + custom_routing_function=Llama4MoE.custom_routing_function, + ) else: - logits = [torch.randn(num_tokens, num_experts, dtype=torch.float32, device=device) for _ in range(5)] - - def kernel_func(): - for router_logits in logits: - moe_layer(hidden_states, topk_layer(hidden_states, router_logits)) - - with benchmark_with_power( - device=device, - kernel_func=kernel_func, - num_warmups=5, - num_runs=10, - repeat_n=1, - ) as results: - pass - return results["latency_ms"] / len(logits), results["power_stats"], kernel_source + topk_layer = TopK( + top_k=topk, + layer_id=0, + use_grouped_topk=num_expert_group is not None and topk_group is not None, + num_expert_group=num_expert_group, + topk_group=topk_group, + renormalize=renormalize, + scoring_func=scoring_func, + correction_bias=correction_bias, + output_format=( + TopKOutputFormat.STANDARD + if routing_method_type is None and moe_backend not in {"flashinfer_mxfp4", "triton_kernel"} + else None + ), + routed_scaling_factor=routed_scaling_factor, + apply_routed_scaling_factor_on_output=( + routed_scaling_factor is not None and moe_layer.should_fuse_routed_scaling_factor_in_topk + ), + is_fp4_experts=is_fp4_experts, + ) + for num_tokens in num_tokens_list: + print(f"num_tokens {num_tokens} ({kernel_source})") + hidden_states = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=device) + if distributed == "balanced": + logits = [ + balanced_logits(num_tokens, num_experts, topk).to(device=device, dtype=torch.float32) + for _ in range(5) + ] + elif distributed == "power_law": + logits = [ + power_law_logits_v3(num_tokens, num_experts, topk, moe_ep_size, power_law_alpha).to( + device=device, dtype=torch.float32 + ) + for _ in range(5) + ] + else: + logits = [ + torch.randn(num_tokens, num_experts, dtype=torch.float32, device=device) for _ in range(5) + ] + + def kernel_func(): + for router_logits in logits: + moe_layer(hidden_states, topk_layer(hidden_states, router_logits)) + + # FlashInfer-backed runners (flashinfer_trtllm / _mxfp4 / _cutlass / _cutedsl) + # are timed on the autotuned tactic like serving; see helper.flashinfer_autotuned. + with ( + flashinfer_autotuned(num_tokens, kernel_func, label=kernel_source), + benchmark_with_power( + device=device, + kernel_func=kernel_func, + num_warmups=5, + num_runs=10, + repeat_n=1, + ) as results, + ): + pass + on_result(num_tokens, results["latency_ms"] / len(logits), results["power_stats"], kernel_source) + return kernel_source finally: if moe_layer is not None: for parameter in moe_layer.parameters(): @@ -1299,10 +1343,57 @@ def _raise_if_unverified_moe_lane(moe_type: str) -> str: _check_compat(specifier, installed_version) for specifier in ("sglang>=0.5.14,<0.5.15", "sglang>=0.5.17,<0.5.18") ) + # 0.5.21 (tag commit e00930c5489053f26d86b179cee0d087f846acbb) re-verified + # for the Hopper/Ada lanes only (2026-10-04): + # * int4_wo: compressed_tensors.py:838-877 — CompressedTensorsMxInt4MoE + # only under an explicit flashinfer_trtllm runner, the Triton WNA16 MoE + # on SM100/103 auto or explicit triton, Marlin for every other CUDA + # path -> SM89/90 stay on Marlin (declared 89/90: marlin). SM100 auto + # moved from flashinfer_trtllm to Triton, so the SM100/103 declaration + # is NOT verified on this series. + # * w4a16_mxfp4 (GPT-OSS): no gpt-oss runner override exists in + # arg_groups/overrides.py any more (model_hook.py:524-528 still points + # at a `_gpt_oss_overrides` that is gone); Mxfp4MoEMethod resolves + # auto in create_moe_runner (mxfp4.py:1423-1430) to TRITON because + # use_triton_kernels is only set by an explicit triton_kernel backend + # (:382) -> SM89/90 declared `triton` matches serving. The SM100/103 + # flashinfer_mxfp4 declaration that 0.5.17 derived from the removed + # override is NOT verified on this series. + # * w4a8_mxfp4_mxfp8 (DSV4 FP4 experts, SM100-only): not re-verified. + if not verified and _check_compat("sglang>=0.5.21,<0.5.22", installed_version): + verified = moe_type in ("int4_wo", "w4a16_mxfp4") and get_sm_version() in (89, 90) + # SM100/103 w4a16_mxfp4 re-verified on 0.5.21 from the B200 identity probes (2026-10-04, + # benchmarks/evidence/collector/op-probe-2026-10/: results/sm100/sglang-0.5.21.yaml + # + results/retests/sm100/sglang-0.5.21.moe_auto.yaml): + # openai/gpt-oss-120b and gpt-oss-20b serve with Mxfp4MoEMethod -> MoeRunnerBackend.FLASHINFER_MXFP4 + # (kernel family trtllm_gen_moe) under the default render AND under an explicit AUTO, i.e. the + # flashinfer_mxfp4 declaration the cases carry for SM100/103 still matches serving even though the + # gpt-oss server-args override is gone (the selection now lives in Mxfp4MoEMethod itself). + if not verified and moe_type == "w4a16_mxfp4" and get_sm_version() in (100, 103): + verified = True + # w4a8_mxfp4_mxfp8 (DeepSeek-V4 native fp4 experts, SM100-only lane) re-verified the same way: the five + # deepseek-ai/DeepSeek-V4* checkpoints serve on B200 with MoeRunnerBackend.FLASHINFER_MXFP4 (trtllm_gen_moe) + # under AUTO (results/retests/sm100/sglang-0.5.21.moe_auto.yaml); the harness's marlin prescription is the + # SM90 workaround ("fp4-experts on SM90: runner=auto crashes"), not the Blackwell path. + if not verified and moe_type == "w4a8_mxfp4_mxfp8" and get_sm_version() in (100, 103): + verified = True + # int4_wo on SM100/103: the source reading above ("SM100 auto moved to Triton") is contradicted + # by the B200 identity records of the same day, which the guard update (fd07bf6b, done on SM90) + # never consumed: moonshotai/Kimi-K2.5 (compressed-tensors W4A16) serves on B200 / 0.5.21 as + # CompressedTensorsFusedMoE -> trtllm_gen_moe under the default render + # (results/sm100/sglang-0.5.21.yaml) AND as moe_runner FLASHINFER_TRTLLM under an explicit AUTO + # (results/retests/sm100/sglang-0.5.21.moe_auto.yaml), i.e. the flashinfer_trtllm declaration the + # SM100/103 cases carry still matches serving. Hardware evidence outranks the source claim + # (harness rule: no claim from reading code). Re-opened 2026-10-06 after the b200_sxm shard run + # (GitLab job 469988017) rejected all 3,078 int4_wo cases on this guard alone. + if not verified and moe_type == "int4_wo" and get_sm_version() in (100, 103): + verified = True if not verified: raise RuntimeError( f"SGLang {moe_type} collection is verified only for the 0.5.14 and 0.5.17 series " - f"(installed: {installed_version}); re-verify framework dispatch before extending this guard." + f"(and 0.5.21 on SM89/SM90/SM100/SM103 for int4_wo / w4a16_mxfp4, SM100/SM103 for " + f"w4a8_mxfp4_mxfp8); installed: {installed_version}, SM{get_sm_version()}; " + "re-verify framework dispatch before extending this guard." ) return installed_version @@ -1314,7 +1405,7 @@ def _triton_config_namespaces() -> list: def run_moe_torch( moe_type, - num_tokens, + num_tokens_list, hidden_size, inter_size, topk, @@ -1346,6 +1437,11 @@ def run_moe_torch( ): torch.cuda.set_device(device) torch.set_default_device(device) + # One case carries the whole token sweep; a single int is accepted for fills and + # direct callers that re-measure one point. + num_tokens_list = [int(num_tokens_list)] if isinstance(num_tokens_list, int) else [int(n) for n in num_tokens_list] + if not num_tokens_list: + raise ValueError("SGLang MoE case has an empty num_tokens_list") assert moe_type in [ "fp8_block", @@ -1387,15 +1483,31 @@ def run_moe_torch( and moe_backend == "flashinfer_cutlass" and (hidden_size % 8 != 0 or local_inter_size % 8 != 0) ): + # FlashInfer 0.6.18 (pinned by SGLang 0.5.21, python/pyproject.toml:42) + # flashinfer_cutlass_fused_moe_binding.cu:386-394: min_alignment = + # 128 / (8 * min(act_bytes, weight_bytes)) = 8 for BF16, ICHECKed on both + # hidden_size and inter_size; SGLang does not pad this path (unquant.py:648-711). raise ValueError( - "SGLang FlashInfer CUTLASS BF16 MoE requires hidden_size and local_inter_size " - f"to be divisible by 8, got hidden_size={hidden_size}, local_inter_size={local_inter_size}" + "FIXME(kernel-limit): SGLang FlashInfer CUTLASS BF16 MoE requires hidden_size and local_inter_size " + "to be divisible by 8 (flashinfer 0.6.18 flashinfer_cutlass_fused_moe_binding.cu:386-394), " + f"got hidden_size={hidden_size}, local_inter_size={local_inter_size}" ) activation_vector_size = 16 if sm_version >= 100 else 8 - if moe_backend == "triton" and is_gated and activation == "gelu" and local_inter_size % activation_vector_size != 0: + if ( + moe_backend == "triton" + and is_gated + and activation in ("gelu", "silu") + and local_inter_size % activation_vector_size != 0 + ): + # sgl-kernel JIT silu_and_mul / gelu_and_mul: activation.cuh:186 + # RuntimeCheck(hidden_size % kVecSize == 0, "hidden size must be divisible + # by vector size"), kVecSize = kMaxVecBytes / sizeof(T) (16 bf16 on SM100+, + # 8 below). 2026-10-11: silu hit it on b200/gb300 (1,053 each) and rtxpro + # (1,377) at local_inter_size 168/84/40 (Nemotron-3-Super inter 2688 at TP16+). raise ValueError( - "SGLang Triton gated BF16 GELU requires local_inter_size " - f"to be divisible by {activation_vector_size}, got local_inter_size={local_inter_size}" + "FIXME(kernel-limit): SGLang Triton gated BF16 GELU/SiLU (sgl-kernel act_and_mul, " + f"activation.cuh:186) requires local_inter_size to be divisible by {activation_vector_size}, " + f"got local_inter_size={local_inter_size} (activation={activation})" ) if ( moe_type == "w4a16_mxfp4" @@ -1404,9 +1516,14 @@ def run_moe_torch( and sm_version == 90 and (hidden_size % 128 != 0 or local_inter_size % 128 != 0) ): + # SGLang 0.5.21 selects Mxfp4FlashinferCutlassMoEMethod on SM90 (fp8.py:448-455), + # which raises unless both widths are multiples of 128 + # (mxfp4_flashinfer_cutlass_moe.py:67); the 256 round-up in layer.py:437-443 + # applies only to the serialized-mxfp4 quant config, not to these FP4 experts. raise ValueError( - "SGLang SM90 DeepSeek-V4 W4A16 FP4 experts require hidden_size and local_inter_size " - f"to be divisible by 128, got hidden_size={hidden_size}, local_inter_size={local_inter_size}" + "FIXME(kernel-limit): SGLang SM90 DeepSeek-V4 W4A16 FP4 experts require hidden_size and local_inter_size " + "to be divisible by 128 (mxfp4_flashinfer_cutlass_moe.py:67 @0.5.21), " + f"got hidden_size={hidden_size}, local_inter_size={local_inter_size}" ) if ( moe_type == "w4a8_mxfp4_mxfp8" @@ -1428,6 +1545,7 @@ def run_moe_torch( # runtime series. (The separate serialized-MXFP4 Mxfp4MoEMethod does # pad, but is not the is_fp4_experts path guarded here.) raise ValueError( + "FIXME(kernel-limit): " "SGLang SM100/103 DeepSeek-V4 W4A8 FP4 experts require hidden_size and local_inter_size " f"to be divisible by 128, got hidden_size={hidden_size}, local_inter_size={local_inter_size}" ) @@ -1458,6 +1576,14 @@ def run_moe_torch( int4_config = get_moe_quantization_module_config("sglang", moe_type, model_name=model_name) int4_group_size = int(int4_config.get("group_size", 128)) if hidden_size % int4_group_size != 0 or local_inter_size % int4_group_size != 0: + # FIXME(re-verify): no 0.5.21 source line rejects K % group_size. The kernels' + # own limits are stricter and fire on their own: SM100/103 MxInt4 packs + # block_k=128 bytes (compressed_tensors_w4a4_mxint4_moe.py:248-254 -> flashinfer + # 0.6.18 fused_moe/core.py:581, hidden/local_inter % 256), SM90 Marlin tiles + # (moe_wna16_marlin.cuh:140-152,265,747: hidden % 128, local_inter % 64). Marlin + # floor-divides the groups (compressed_tensors_wNa16_moe.py:158-159), so an + # unaligned K may run with a different group size than the row says — not shown + # either way, so this stays an unclassified failure until a run decides it. raise ValueError( "SGLang INT4-WO group quantization requires hidden_size and local_inter_size " f"to be divisible by group_size={int4_group_size}, got " @@ -1466,9 +1592,14 @@ def run_moe_torch( block_shape = [0, int4_group_size] elif moe_type == "fp8_block": if moe_backend == "triton" and (hidden_size % 128 != 0 or local_inter_size % 128 != 0): + # Both GEMM inputs are per-token-group quantized at block_k=128 + # (fused_moe_triton_kernels.py:866 -> fp8_kernel.py:577 + # `assert x.shape[-1] % group_size == 0` @0.5.21): K = hidden_size for the + # first GEMM, local_inter_size for the second. raise ValueError( - "SGLang Triton fp8_block requires hidden_size and local_inter_size " - f"to be divisible by 128, got hidden_size={hidden_size}, local_inter_size={local_inter_size}" + "FIXME(kernel-limit): SGLang Triton fp8_block requires hidden_size and local_inter_size " + "to be divisible by 128 (fp8_kernel.py:577 via fused_moe_triton_kernels.py:866 @0.5.21), " + f"got hidden_size={hidden_size}, local_inter_size={local_inter_size}" ) block_shape = [128, 128] else: @@ -1486,15 +1617,68 @@ def run_moe_torch( or apply_router_weight_on_input or (moe_backend != "triton" and not use_int4_w4a16) ) - # Record (never choose) whether the timed Triton call used a tuned JSON or + distribution = "power_law_" + str(power_law_alpha) if distributed == "power_law" else distributed + + # Record (never choose) whether each timed Triton call used a tuned JSON or # SGLang's default config; both the framework runner and the direct path - # resolve it through the same two config functions. + # resolve it through the same two config functions. One task times a whole + # num_tokens series, so the counters cover one point at a time: _persist + # hands them to the sidecar with that point's row and resets them. with observe_triton_config_selection(_triton_config_namespaces()) as triton_config_observed: + + def _persist(num_tokens: int, latency: float, power_stats: dict, kernel_source: str) -> None: + persisted = log_perf( + item_list=[ + { + "moe_dtype": moe_type, + "num_tokens": num_tokens, + "hidden_size": hidden_size, + "inter_size": inter_size, + "topk": topk, + "num_experts": num_experts, + "moe_tp_size": moe_tp_size, + "moe_ep_size": moe_ep_size, + "distribution": distribution, + "latency": latency, + } + ], + framework="SGLang", + version=installed_sglang_version, + device_name=torch.cuda.get_device_name(device), + op_name="moe", + kernel_source=kernel_source, + perf_filename=perf_filename, + power_stats=power_stats, + ) + if not persisted: + raise RuntimeError("Failed to persist SGLang MoE performance row") + point_observed = dict(triton_config_observed) + for key in triton_config_observed: + triton_config_observed[key] = 0 + if moe_backend == "triton": + append_triton_config_record( + perf_filename, + { + "moe_dtype": moe_type, + "num_tokens": num_tokens, + "hidden_size": hidden_size, + "inter_size": inter_size, + "topk": topk, + "num_experts": num_experts, + "moe_tp_size": moe_tp_size, + "moe_ep_size": moe_ep_size, + "device_name": torch.cuda.get_device_name(device), + }, + kernel_source, + point_observed, + ) + if use_framework_layer: - latency, power_stats, kernel_source = _benchmark_framework_quantized_moe( + _benchmark_framework_quantized_moe( moe_type=moe_type, moe_backend=moe_backend, - num_tokens=num_tokens, + num_tokens_list=num_tokens_list, + on_result=_persist, hidden_size=hidden_size, inter_size=inter_size, topk=topk, @@ -1523,103 +1707,60 @@ def run_moe_torch( device=device, ) else: - rank0_workloads: list[Rank0Workload] | None = None - if moe_ep_size > 1 and distributed in ("power_law", "balanced"): - rank0_workloads = build_rank0_workloads( - num_workloads=5, - num_tokens=num_tokens, - hidden_size=hidden_size, - topk=topk, - num_experts=num_experts, - moe_ep_size=moe_ep_size, - distributed=distributed, - power_law_alpha=power_law_alpha if distributed == "power_law" else None, - dtype=torch.bfloat16, - device=torch.device(device), - ) + kernel_source = { + "triton": "sglang_fused_moe_triton", + "triton_kernel": "sglang_triton_kernel_moe", + "marlin": "sglang_marlin_moe", + "flashinfer_trtllm": "sglang_flashinfer_trtllm_moe", + "flashinfer_cutlass": "sglang_flashinfer_cutlass_moe", + }[moe_backend] try: - latency, power_stats = benchmark( - num_tokens, - num_local_experts, - (2 if is_gated else 1) * local_inter_size, - hidden_size, - topk, - torch.bfloat16, - moe_type == "fp8_block", - False, - False, - block_shape=block_shape, - distributed=distributed, - power_law_alpha=power_law_alpha, - workloads=rank0_workloads, - swiglu_limit=swiglu_limit, - activation=activation, - is_gated=is_gated, - gemm1_alpha=gemm1_alpha, - gemm1_clamp_limit=gemm1_clamp_limit, - ) + for num_tokens in num_tokens_list: + rank0_workloads: list[Rank0Workload] | None = None + if moe_ep_size > 1 and distributed in ("power_law", "balanced"): + rank0_workloads = build_rank0_workloads( + num_workloads=5, + num_tokens=num_tokens, + hidden_size=hidden_size, + topk=topk, + num_experts=num_experts, + moe_ep_size=moe_ep_size, + distributed=distributed, + power_law_alpha=power_law_alpha if distributed == "power_law" else None, + dtype=torch.bfloat16, + device=torch.device(device), + ) + latency, power_stats = benchmark( + num_tokens, + num_local_experts, + (2 if is_gated else 1) * local_inter_size, + hidden_size, + topk, + torch.bfloat16, + moe_type == "fp8_block", + False, + False, + block_shape=block_shape, + distributed=distributed, + power_law_alpha=power_law_alpha, + workloads=rank0_workloads, + swiglu_limit=swiglu_limit, + activation=activation, + is_gated=is_gated, + gemm1_alpha=gemm1_alpha, + gemm1_clamp_limit=gemm1_clamp_limit, + ) + _persist(num_tokens, latency, power_stats, kernel_source) except Exception: _fmoe_kernels_mod._B_DESC_CACHE.clear() gc.collect() torch.cuda.empty_cache() raise - - if not use_framework_layer: - _fmoe_kernels_mod._B_DESC_CACHE.clear() - gc.collect() - torch.cuda.empty_cache() + _fmoe_kernels_mod._B_DESC_CACHE.clear() + gc.collect() + torch.cuda.empty_cache() restart_worker = torch.cuda.memory_allocated(device) > torch.cuda.get_device_properties(device).total_memory // 4 - if not use_framework_layer: - kernel_source = { - "triton": "sglang_fused_moe_triton", - "triton_kernel": "sglang_triton_kernel_moe", - "marlin": "sglang_marlin_moe", - "flashinfer_trtllm": "sglang_flashinfer_trtllm_moe", - "flashinfer_cutlass": "sglang_flashinfer_cutlass_moe", - }[moe_backend] - persisted = log_perf( - item_list=[ - { - "moe_dtype": moe_type, - "num_tokens": num_tokens, - "hidden_size": hidden_size, - "inter_size": inter_size, - "topk": topk, - "num_experts": num_experts, - "moe_tp_size": moe_tp_size, - "moe_ep_size": moe_ep_size, - "distribution": "power_law_" + str(power_law_alpha) if distributed == "power_law" else distributed, - "latency": latency, - } - ], - framework="SGLang", - version=installed_sglang_version, - device_name=torch.cuda.get_device_name(device), - op_name="moe", - kernel_source=kernel_source, - perf_filename=perf_filename, - power_stats=power_stats, - ) - if not persisted: - raise RuntimeError("Failed to persist SGLang MoE performance row") - if moe_backend == "triton": - append_triton_config_record( - perf_filename, - { - "moe_dtype": moe_type, - "num_tokens": num_tokens, - "hidden_size": hidden_size, - "inter_size": inter_size, - "topk": topk, - "num_experts": num_experts, - "moe_tp_size": moe_tp_size, - "moe_ep_size": moe_ep_size, - "device_name": torch.cuda.get_device_name(device), - }, - kernel_source, - triton_config_observed, - ) if restart_worker: return WORKER_RESTART diff --git a/python/aisimulate/collector/sglang/collect_msa_module.py b/python/aisimulate/collector/sglang/collect_msa_module.py index 30f66baa17..a9a495a1ca 100644 --- a/python/aisimulate/collector/sglang/collect_msa_module.py +++ b/python/aisimulate/collector/sglang/collect_msa_module.py @@ -2,10 +2,10 @@ # SPDX-License-Identifier: Apache-2.0 # MiniMax-M3 support landed in SGLang v0.5.16 (models/minimax_m3.py + -# layers/attention/minimax_sparse_backend.py + minimax_sparse_ops/); the -# collector-default 0.5.14 pin has no M3. The framework manifest pins the -# msa family to the official v0.5.16 image. -__compat__ = "sglang>=0.5.16" +# layers/attention/minimax_sparse_backend.py + minimax_sparse_ops/); the msa +# family now runs on the 0.5.21 default (the v0.5.16 family pin is retired). +# One release per file (owner policy 2026-09-20); earlier code is `git log -- ` away. +__compat__ = "sglang==0.5.21" """ MSA Module Collector for SGLang — MiniMax-M3 sparse-attention benchmarking. @@ -78,10 +78,13 @@ get_mla_module_sweep_spec, ) from collector.helper import ( + PhaseClock, _resolve_local_model_path, benchmark_with_power, get_sm_version, log_perf, + render_phase_totals, + run_measurement_subprocess, ) try: @@ -136,6 +139,10 @@ # All three are emulated on one GPU by shrinking the config head counts. +# Per-subprocess sum of the point phase times printed by run_msa_module. +_POINT_PHASE_TOTALS: dict[str, float] = {} + + def _ceil_div(a: int, b: int) -> int: return (a + b - 1) // b @@ -323,17 +330,31 @@ def _generation_shapes(max_pos: int | None): _M3_NATIVE_GQA_RATIO = 16 # 64 q / 4 kv heads (bundled config) +def _loaded_num_layers() -> int: + """Layers the worker loads (SGLANG_TEST_NUM_LAYERS, default 2) — every one of + them owns a slot in the main pool and the index-K pool, so the plan-time + footprint is per token x layers (memory_pool.py:5404 MiniMaxSparseKVPool + builds one MHA pool over start_layer..end_layer plus one index pool per + sparse layer @0.5.21).""" + return int(os.environ.get("SGLANG_TEST_NUM_LAYERS", "2")) + + def _estimated_kv_bytes_per_token(num_heads: int, kv_cache_dtype: str) -> int: - """Upper-bound bytes/token of the M3 pools at this per-GPU head count: - main paged K+V (kv-head sharded; fp8 stores 1 B/elem, else bf16) plus the - index-K side cache (always model dtype, kv_cache_configurator.py:1246 - index_dtype=model_dtype @v0.5.16).""" + """Upper-bound bytes/token of the M3 pools at this per-GPU head count, for + ALL loaded layers: main paged K+V (kv-head sharded; fp8 stores 1 B/elem, else + bf16) plus the index-K side cache (model dtype unless fp8 attn-GEMM, + memory_pool.py:5392 get_minimax_sparse_index_dtype @0.5.21; index_dtype= + model_dtype @v0.5.16 kv_cache_configurator.py:1246). The 0.5.16-era + estimate priced one layer while the worker loads two: the three + largest-batch decode cells (b=256/512/1024 x 131072/65536/32768) passed the + plan filter and then failed the worker's execute-or-raise capacity check + (33.6M planned tokens > 24.5M real pool, H20 0.5.21 smoke 2026-10-01).""" kv_heads = max(1, num_heads // _M3_NATIVE_GQA_RATIO) main_elem = 1 if kv_cache_dtype == "fp8" else 2 main = 2 * kv_heads * _M3_HEAD_DIM * main_elem index_heads = min(_M3_INDEX_HEADS, num_heads) index = index_heads * _M3_INDEX_DIM * 2 - return main + index + return (main + index) * _loaded_num_layers() def _plan_memory_filter(shapes, *, num_heads: int, kv_cache_dtype: str, is_prefill: bool): @@ -400,11 +421,12 @@ def _build_module_test_cases(mode: str): empty_groups = 0 for model_spec in model_specs: for compute_dtype, kv_dtype, gemm_type in _get_precision_combos(mode): - for target_tp in sweep.module_tp_sizes: + # the model row's declared shards win over the sweep's module_tp_sizes + for target_tp in model_spec.tensor_parallel_sizes or sweep.module_tp_sizes: if model_spec.native_num_heads % target_tp != 0: continue num_heads = model_spec.native_num_heads // target_tp - if num_heads not in sweep.inner_sweep_head_counts: + if model_spec.tensor_parallel_sizes is None and num_heads not in sweep.inner_sweep_head_counts: continue batch_sizes = sweep.context_batch_sizes if mode == "context" else [0] for batch_size in batch_sizes: @@ -656,11 +678,7 @@ def load_model_runner( (server_args.py _handle_gpu_memory_settings) and would go negative. """ from sglang.srt.configs.model_config import ModelConfig - from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.entrypoints.engine import _set_envs_and_config - from sglang.srt.layers.moe import initialize_moe_config - from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config - from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.server_args import ServerArgs from sglang.srt.utils import suppress_other_loggers @@ -725,9 +743,6 @@ def load_model_runner( ) _set_envs_and_config(server_args) - initialize_moe_config(server_args) - initialize_fp8_gemm_config(server_args) - initialize_fp4_gemm_config(server_args) model_config = ModelConfig.from_server_args(server_args) actual_architecture = (model_config.hf_config.architectures or [None])[0] @@ -745,12 +760,15 @@ def load_model_runner( _sock.bind(("127.0.0.1", 0)) nccl_port = _sock.getsockname()[1] + from collector.sglang.runtime_compat import init_runtime_config + + _runner_parallel_kwargs = init_runtime_config(server_args, gpu_id, nccl_port=nccl_port, model_config=model_config) model_runner = ModelRunner( model_config=model_config, mem_fraction_static=server_args.mem_fraction_static, gpu_id=gpu_id, - ps=ParallelState.trivial(gpu_id=gpu_id), nccl_port=nccl_port, + **_runner_parallel_kwargs, server_args=server_args, ) model_runner.alloc_memory_pool() @@ -840,7 +858,13 @@ def _decode_graph_covered(model_runner, num_tokens: int) -> bool: graph: cuda_graph_config.decode coverage (backend != disabled; bs list or max_bs — resolved per GPU tier by server_args _handle_gpu_memory_settings, e.g. 256 on the SM90 validation node at tp<4).""" - decode_cfg = model_runner.server_args.cuda_graph_config.decode + cuda_graph_config = model_runner.server_args.cuda_graph_config + if cuda_graph_config is None: + # sglang>=0.5.21: derived config lives behind the resolving view (arg_groups/model_override_base.py:120) + from sglang.srt.arg_groups.model_override_base import resolving_view + + cuda_graph_config = resolving_view(model_runner.server_args).cuda_graph_config + decode_cfg = cuda_graph_config.decode if decode_cfg.backend == "disabled": return False if decode_cfg.bs: @@ -996,6 +1020,7 @@ def _run_prefill_point( """ from sglang.srt.model_executor.forward_batch_info import ForwardBatch + point_clock = PhaseClock() model_runner.req_to_token_pool.clear() model_runner.token_to_kv_pool_allocator.clear() @@ -1010,6 +1035,7 @@ def _run_prefill_point( batch = _make_schedule_batch(model_runner, reqs) with temporarily_chunked_alloc_extend(model_runner, batch_size * seq_len): batch.prepare_for_extend() + point_clock.mark("batch") # return_hidden_states_before_norm=False mirrors the serving worker # (managers/tp_worker.py:267,548@v0.5.16). forward_batch = ForwardBatch.init_new(batch, model_runner, return_hidden_states_before_norm=False) @@ -1019,6 +1045,7 @@ def _run_prefill_point( # minimax_sparse_backend.py:505-507, base_attn_backend.py:48-54). with _forward_ctx(model_runner): model_runner.attn_backend.init_forward_metadata(forward_batch) + point_clock.mark("metadata") hidden_states = torch.randn( batch_size * seq_len, @@ -1033,6 +1060,7 @@ def _run_prefill_point( .contiguous() .flatten() ) + point_clock.mark("inputs") def kernel_func(): with _forward_ctx(model_runner): @@ -1053,11 +1081,16 @@ def kernel_func(): use_cuda_graph=False, ) as results: pass + point_clock.mark("bench") + point_clock.absorb(results.get("phase_seconds"), prefix="bench_") return results["latency_ms"], results["power_stats"] finally: model_runner.req_to_token_pool.clear() model_runner.token_to_kv_pool_allocator.clear() torch.cuda.empty_cache() + point_clock.mark("cleanup") + point_clock.add_to(_POINT_PHASE_TOTALS) + print(f" [MSA point timing] context b={batch_size}, s={seq_len}, prefix={prefix_len}: {point_clock.render()}") def _run_decode_point( @@ -1083,6 +1116,7 @@ def _run_decode_point( """ from sglang.srt.model_executor.forward_batch_info import ForwardBatch + point_clock = PhaseClock() model_runner.req_to_token_pool.clear() model_runner.token_to_kv_pool_allocator.clear() @@ -1094,6 +1128,7 @@ def _run_decode_point( req.output_ids.append(0) batch.prepare_for_decode() forward_batch = ForwardBatch.init_new(batch, model_runner, return_hidden_states_before_norm=False) + point_clock.mark("batch") use_graph = _decode_graph_covered(model_runner, batch_size) with _forward_ctx(model_runner): @@ -1108,6 +1143,7 @@ def _run_decode_point( # the serving decode graph runner owns — the module bench never # executes a dense layer, so its eager metadata is kept. model_runner.attn_backend.sparse.init_forward_metadata_out_graph(forward_batch, in_capture=True) + point_clock.mark("metadata") hidden_states = torch.randn( batch_size, @@ -1116,6 +1152,7 @@ def _run_decode_point( device="cuda", ) positions = torch.full((batch_size,), kv_len, device="cuda", dtype=torch.int64) + point_clock.mark("inputs") def kernel_func(): with _forward_ctx(model_runner): @@ -1136,11 +1173,16 @@ def kernel_func(): use_cuda_graph=use_graph, ) as results: pass + point_clock.mark("bench") + point_clock.absorb(results.get("phase_seconds"), prefix="bench_") return results["latency_ms"], results["power_stats"], use_graph finally: model_runner.req_to_token_pool.clear() model_runner.token_to_kv_pool_allocator.clear() torch.cuda.empty_cache() + point_clock.mark("cleanup") + point_clock.add_to(_POINT_PHASE_TOTALS) + print(f" [MSA point timing] generation b={batch_size}, kv={kv_len}: {point_clock.render()}") # ═══════════════════════════════════════════════════════════════════════ @@ -1342,6 +1384,7 @@ def run_msa_module( summary = f"ok={ok} error={len(failures)} skip=0 total={len(shapes)}" print(f"[MSA] {phase} {summary}") + print(f"[MSA timing] {phase}: {render_phase_totals(_POINT_PHASE_TOTALS)}") if ok == 0: raise RuntimeError(f"MSA module {phase} persisted no rows; {summary}") if failures: @@ -1390,34 +1433,33 @@ def _run_msa_subprocess( f"{batch_repr}, {target_tp_size}, None, {inner_manifest!r})\n" ) - proc = subprocess.Popen( - [sys.executable, "-c", code], - env=env, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - cwd=os.path.dirname(os.path.abspath(__file__)), - ) _timeout_env = os.environ.get("AIC_MLA_MODULE_SUBPROCESS_TIMEOUT_SEC") subprocess_timeout = int(_timeout_env) if _timeout_env else None try: - stdout, _ = proc.communicate(timeout=subprocess_timeout) + # beats the stall watchdog while the child prints progress; the child dies with the worker + returncode, stdout = run_measurement_subprocess( + [sys.executable, "-c", code], + env=env, + cwd=os.path.dirname(os.path.abspath(__file__)), + timeout=subprocess_timeout, + ) if stdout: print(stdout.decode("utf-8", errors="replace")) except subprocess.TimeoutExpired as exc: - proc.kill() - proc.wait() + if exc.output: + print(exc.output.decode("utf-8", errors="replace")) raise RuntimeError( f"MSA module {phase} subprocess timed out after {subprocess_timeout}s " f"(heads={num_heads}, model={model_path}, kv={kv_cache_dtype})" ) from exc - if proc.returncode != 0: + if returncode != 0: tail = "" if stdout: lines = stdout.decode("utf-8", errors="replace").strip().splitlines() tail = "\n".join(lines[-30:]) raise RuntimeError( - f"MSA module {phase} subprocess failed (exit code {proc.returncode})\n" + f"MSA module {phase} subprocess failed (exit code {returncode})\n" f"--- subprocess output (last 30 lines) ---\n{tail}" ) diff --git a/python/aisimulate/collector/sglang/deepseekv4_sparse_modules.py b/python/aisimulate/collector/sglang/deepseekv4_sparse_modules.py index 54083fc9bf..589f3d397f 100644 --- a/python/aisimulate/collector/sglang/deepseekv4_sparse_modules.py +++ b/python/aisimulate/collector/sglang/deepseekv4_sparse_modules.py @@ -5,35 +5,39 @@ Collects the DeepSeek-V4 NSA sparse-attention sub-kernels at the kernel level. DeepSeek-V4's primary AIC path uses prefix-aware full CSA/HCA attention-module -rows; paged_mqa_logits / csa_attn / hca_attn here are supporting kernel-level -data for prefix/past_kv correction and residual analysis, while the topk_calib -DELTA is actively consumed by perf_database's topK correction. +rows; paged_mqa_logits here is supporting kernel-level data for the SDK CP +chunk walk, while the topk_calib DELTA is actively consumed by +perf_database's topK correction. -The four sub-kernels form ONE sparse-op family — the CSA/HCA sparse path in +The two sub-kernels form ONE sparse-op family — the CSA sparse path in order — all collected, none modeled analytically: 1. ``deep_gemm.fp8_paged_mqa_logits`` CSA indexer scoring 2. ``topk_transform`` (topk_512/_1024) CSA selection — phase-specific v1/v2 flat-vs-top_last DELTA calib - 3. ``flash_mla_with_kvcache`` (csa_attn) CSA sparse FMLA over the - topk-selected c4 positions - 4. ``flash_mla_with_kvcache`` (hca_attn) HCA c128 sparse FMLA + +The standalone FMLA sidecars (``flash_mla_with_kvcache`` as csa_attn / hca_attn) +were removed 2026-10-15: their SDK readers were retired design placeholders with +no query path (csa_attn never shipped data) and, as registry-only ops, they +failed every shard plan generated from the registry. Inputs are SAME-SOURCE: every sub-kernel derives its ``(prefix, isl, bs[, tp])`` -shapes STRICTLY 1:1 from the CSA/HCA attention-module CSV it belongs to -(paged_mqa / topk / csa_attn ← CSA module; hca_attn ← HCA module) — no separate -sweep grid — benched at upstream layouts (DeepGEMM ``test_attention.py`` for the -indexer GEMM, FlashMLA ``MODEL1_FP8Sparse`` quant for the FMLA kernel). +shapes STRICTLY 1:1 from the CSA attention-module INPUT sweep it belongs to — no +separate sweep grid — benched at the upstream layout (DeepGEMM +``test_attention.py`` for the indexer GEMM). CSV schema matches existing aic dsv4 module CSVs (so loaders can be shared): -``isl`` carries M, ``step`` carries past_kv, ``compress_ratio`` distinguishes -CSA(=4) / HCA(=128). +``isl`` carries M, ``step`` carries past_kv, ``compress_ratio`` carries the CSA +c4 path (=4). """ -# Requires stock SGLang 0.5.14 with its matching ``sgl-kernel`` package. +# Requires stock SGLang 0.5.21 with its matching ``sgl-kernel`` package. from __future__ import annotations -__compat__ = "sglang==0.5.14" +# 0.5.21 since 2026-10-01 (H20/sm90 collector port: op_smoke + path gates in the v0.5.21 image; +# findings hopper_sglang_collector_port_0514_to_0521_2026_10_01). +# One release per file (owner policy 2026-09-20); earlier code is `git log -- ` away. +__compat__ = "sglang==0.5.21" import functools import json @@ -63,10 +67,10 @@ def _dsv4_sparse_kernel_cases(kernel, models=None): # One task per (model, bs); collect.py spreads bs across GPU workers (no - # single-worker cuda-graph buildup). All sparse kernels run a single fixed - # head config: the FMLA pads to its required head count OUTSIDE the kernel - # (model TP zero-pad), so the kernel is TP-independent -> one config per bs, - # no tp sweep. Each task sweeps (isl, prefix) for its bs. + # single-worker cuda-graph buildup). The sparse kernels run a single fixed + # head config (the indexer heads are not TP-sharded), so each is + # TP-independent -> one config per bs, no tp sweep. Each task sweeps + # (isl, prefix) for its bs. if models is None: try: from collector.case_generator import _selected_dsv4_models @@ -87,24 +91,12 @@ def get_dsv4_paged_mqa_logits_test_cases(): return _dsv4_sparse_kernel_cases("paged_mqa_logits") -def get_dsv4_hca_attn_test_cases(): - return _dsv4_sparse_kernel_cases("hca_attn") - - -def get_dsv4_csa_attn_test_cases(): - return _dsv4_sparse_kernel_cases("csa_attn") - - get_dsv4_flash_paged_mqa_logits_test_cases = get_dsv4_paged_mqa_logits_test_cases -get_dsv4_flash_hca_attn_test_cases = get_dsv4_hca_attn_test_cases __all__ = [ "DEFAULT_MODEL", - "get_dsv4_csa_attn_test_cases", - "get_dsv4_flash_hca_attn_test_cases", "get_dsv4_flash_paged_mqa_logits_test_cases", - "get_dsv4_hca_attn_test_cases", "get_dsv4_paged_mqa_logits_test_cases", "get_dsv4_topk_calib_test_cases", "run_dsv4_sparse_kernel_worker", @@ -180,25 +172,6 @@ def _dsv4_sparse_config(model_path: str) -> SimpleNamespace: ) -# ═══════════════════════════════════════════════════════════════════════ -# DeepSeek-V4 sparse architectural constants -# ═══════════════════════════════════════════════════════════════════════ -# Shape fields below are identical across DSV4 variants (Flash/Pro), so the -# default model config is authoritative; read them via the sparse config. -_DEFAULT_SC = _dsv4_sparse_config(DEFAULT_MODEL) -V_HEAD_DIM = _DEFAULT_SC.head_dim # 512 -FMLA_D_QK = V_HEAD_DIM # FlashMLA d_qk == head_dim (NOT 576) -FMLA_D_ROPE = _DEFAULT_SC.d_rope # 64 -FMLA_D_NOPE = _DEFAULT_SC.d_nope # 448 - -# Kernel-layout invariants — NOT in the model config. The FlashMLA -# MODEL1_FP8Sparse fp8 quant tile size and the deep_gemm / FlashMLA paged block -# size are hardcoded in the kernels themselves (deep_gemm blocksize 64; -# flashmla_backend.PAGE_SIZE = 64), so they are fixed constants here. -# bytes_per_token = d_nope + d_rope*2 + num_tiles + 1 pad = 448 + 128 + 7 + 1 = 584 -FMLA_TILE_SIZE = 64 -FMLA_NUM_TILES = FMLA_D_NOPE // FMLA_TILE_SIZE # 448 / 64 = 7 - DEFAULT_ARCHITECTURE = "DeepseekV4ForCausalLM" @@ -224,49 +197,36 @@ def _device_num_sms(device: str | torch.device) -> int: return torch.cuda.get_device_properties(device).multi_processor_count -# The DSV4 sparse-op family — all four collected by ONE worker +# The DSV4 sparse-op family — both collected by ONE worker # (run_dsv4_sparse_kernel_worker), all driven by the KERNEL_* maps below, all -# same-source (shapes read 1:1 from the CSA/HCA attention-module CSV they belong -# to), all writing the SAME row schema (one ``latency`` per row): +# same-source (shapes derived 1:1 from the CSA attention-module INPUT sweep), +# all writing the SAME row schema (one ``latency`` per row): # - paged_mqa_logits: CSA indexer scoring (deep_gemm) -> 1 row/shape # - topk: CSA selection — context runs v1 and generation runs v2, # each under two score distributions, emitted as TWO # rows/shape (score_mode=v1_flat | v1_top_last or # v2_flat | v2_top_last); perf_database takes the matching # DELTA = flat.latency - top_last.latency -> 2 rows/shape -# - csa_attn: CSA's flash_mla over the topk-selected c4 -# positions (K_per_query = min(index_topk, full_s//4)) -# - hca_attn: HCA's flash_mla over c128 cache. -# (production DSV4 at TP>1 also runs FlashMLA with -# the native h_q — sglang pads Q to full native heads, -# then slices output back; see deepseek_v4.py:847. So -# TP=1 data is valid for any deployment TP.) KERNEL_TO_OP_NAME = { "paged_mqa_logits": "dsv4_paged_mqa_logits_module", "topk": "dsv4_csa_topk_calib", - "hca_attn": "dsv4_hca_attn_module", - "csa_attn": "dsv4_csa_attn_module", } KERNEL_TO_KERNEL_SOURCE = { "paged_mqa_logits": "deep_gemm.fp8_paged_mqa_logits", - "hca_attn": "flash_mla_with_kvcache", - "csa_attn": "flash_mla_with_kvcache", } -# compress_ratio: 4 for the CSA c4 path (indexer/topk/csa_attn), 128 for HCA c128. +# compress_ratio: 4 for the CSA c4 path (indexer/topk). KERNEL_TO_COMPRESS_RATIO = { "paged_mqa_logits": 4, "topk": 4, - "hca_attn": 128, - "csa_attn": 4, } # Each sparse sub-kernel derives its (prefix, isl, bs) shapes from its owning -# attention module's INPUT sweep (paged_mqa/topk/csa_attn <- CSA, hca_attn <- -# HCA) via _dsv4_context_derived_shapes / _dsv4_generation_derived_shapes — -# NOT from the already-collected module CSVs. -# ALL four sparse kernels are TP-independent — the indexer/FMLA latency does not +# attention module's INPUT sweep (paged_mqa/topk <- CSA) via +# _dsv4_context_derived_shapes / _dsv4_generation_derived_shapes — NOT from the +# already-collected module CSVs. +# Both sparse kernels are TP-independent — the indexer latency does not # change with tp (the module's per-tp rows differ only by the attention module's # comms/sharding, not the sparse sub-kernel). So each is benched ONCE per unique # (prefix, isl, bs) module shape and written one row per shape (tp_size=1), @@ -326,7 +286,7 @@ def timed_kernel(): # ═══════════════════════════════════════════════════════════════════════ -# Cache packing helpers (mirror DeepGEMM/FlashMLA test code) +# Cache packing helpers (mirror DeepGEMM test code) # ═══════════════════════════════════════════════════════════════════════ @@ -348,65 +308,6 @@ def _kv_cache_cast_to_fp8_indexer(x: torch.Tensor) -> torch.Tensor: return out.view(num_blocks, block_size, 1, head_dim + 4) -def _quantize_k_cache_model1(k_bf16: torch.Tensor) -> torch.Tensor: - """FlashMLA MODEL1_FP8Sparse pack (DSV4 sparse layout). - - k_bf16: (num_blocks, block_size, 1, d_qk=512) bf16 - out: (num_blocks, block_size, 1, bytes_per_token=584) packed fp8 - - Layout per token (inside the bytes_per_token slab): - [d_nope=448 fp8 nope][d_rope*2=128 bf16 rope][num_tiles=7 fp8 e8m0 sf][1 pad] - """ - num_blocks, block_size, num_heads, d_qk = k_bf16.shape - assert num_heads == 1 and d_qk == FMLA_D_QK - k = k_bf16.squeeze(2) # (num_blocks, block_size, d_qk) - - bytes_per_token = FMLA_D_NOPE + 2 * FMLA_D_ROPE + FMLA_NUM_TILES + 1 # 584 - size_per_block_padded = (block_size * bytes_per_token + 576 - 1) // 576 * 576 - # Allocate padded (so memory layout matches kernel TMA expectations) but - # SLICE back to exact ``block_size*bytes_per_token`` so the final view has - # last-dim = bytes_per_token (kernel asserts on this). - out = torch.empty((num_blocks, size_per_block_padded), dtype=torch.float8_e4m3fn, device=k_bf16.device) - out_view = out[:, : block_size * bytes_per_token] - - # nope+rope region - nope_rope_part = out_view[:, : block_size * (FMLA_D_NOPE + 2 * FMLA_D_ROPE)].view( - num_blocks, block_size, FMLA_D_NOPE + 2 * FMLA_D_ROPE - ) - nope = nope_rope_part[:, :, :FMLA_D_NOPE] - rope = nope_rope_part[:, :, FMLA_D_NOPE:].view(k_bf16.dtype) - - sf = ( - out_view[:, block_size * (FMLA_D_NOPE + 2 * FMLA_D_ROPE) :] - .view(num_blocks, block_size, 8)[:, :, :7] - .view(torch.float8_e8m0fnu) - ) - - rope[:] = k[..., FMLA_D_NOPE:] - for tile_idx in range(FMLA_NUM_TILES): - s, e = tile_idx * FMLA_TILE_SIZE, (tile_idx + 1) * FMLA_TILE_SIZE - scale_inv = (k[..., s:e].abs().float().amax(dim=-1).float() / 448.0).clamp_min(1e-4) - scale_inv = torch.pow(2, scale_inv.log2().ceil()).to(torch.float32) - sf[:, :, tile_idx] = scale_inv.to(torch.float8_e8m0fnu) - - scale_inv = scale_inv.unsqueeze(-1) - nope[:, :, s:e] = (k[..., s:e].float() / scale_inv.float()).to(torch.float8_e4m3fn) - - # Return a view whose stride(0) is the padded per-block size — FlashMLA's - # MODEL1 path asserts ``k_cache.stride(0) % TMA_K_STRIDE == 0`` (with - # TMA_K_STRIDE = D_NOPE + 2*D_ROPE = 576), and ``bytes_per_token * block_size`` - # = 37376 is *not* a multiple of 576, so we need the per-block padding to be - # visible in the tensor's stride. ``.view`` collapses stride(0) to the - # contiguous value when num_blocks == 1, breaking the assertion for small - # shapes; ``as_strided`` lets us pin stride(0) to size_per_block_padded for - # any num_blocks. - return out.as_strided( - size=(num_blocks, block_size, 1, bytes_per_token), - stride=(size_per_block_padded, bytes_per_token, bytes_per_token, 1), - storage_offset=0, - ) - - def _expand_block_table(num_blocks: int, batch_size: int, device) -> torch.Tensor: """Block list [0..num_blocks) replicated to (batch_size, num_blocks) int32. All requests share the same physical blocks (KV cache is shared in-bench).""" @@ -414,12 +315,6 @@ def _expand_block_table(num_blocks: int, batch_size: int, device) -> torch.Tenso return row.unsqueeze(0).expand(batch_size, num_blocks).contiguous() -def _expand_indices(k: int, batch_size: int, m: int, device) -> torch.Tensor: - """Sequential KV indices [0..k) replicated to (batch_size, m, k) int32.""" - base = torch.arange(k, dtype=torch.int32, device=device) - return base.view(1, 1, k).expand(batch_size, m, k).contiguous() - - # ═══════════════════════════════════════════════════════════════════════ # Kernel 1: deep_gemm.fp8_paged_mqa_logits # ═══════════════════════════════════════════════════════════════════════ @@ -440,7 +335,19 @@ def _bench_paged_mqa_logits( to ``batch_dim`` with ``next_n=1``. Per-request causal lengths remain distinct even though those token rows share one physical KV cache. """ - from deep_gemm import fp8_paged_mqa_logits, get_paged_mqa_logits_metadata + from deep_gemm import fp8_paged_mqa_logits + + try: + # sglang>=0.5.21 builds the paged-MQA schedule with its own JIT kernel + # (sglang::paged_mqa_metadata_tiny_kernel) whenever SGLANG_OPT_USE_JIT_INDEXER_METADATA + # is set (default True, environ.py:1570) or the query row count exceeds 11673 + # (dsv4/metadata.py:213-224, force_deep_gemm_metadata defaults False :184); the + # DeepGEMM sched kernel is only the opt-out path. Same (seq_lens, page_size, num_sm) + # contract (kernels/ops/attention/dsv4/attn.py:57). H20 gate 2026-10-01: the serving + # decode trace carries the tiny kernel, never deep_gemm::sched::sm90_paged_mqa_logits_metadata. + from sglang.kernels.ops.attention.dsv4.attn import get_paged_mqa_logits_metadata + except ImportError: + from deep_gemm import get_paged_mqa_logits_metadata if batch_size <= 0 or M % batch_size: raise ValueError(f"M={M} must be divisible by positive batch_size={batch_size}") @@ -481,173 +388,6 @@ def kernel_fn(): return _bench_cuda_graph(kernel_fn, allow_graph_fail=False, device=device) -# ═══════════════════════════════════════════════════════════════════════ -# Helpers for FlashMLA HCA -# ═══════════════════════════════════════════════════════════════════════ - - -def _build_flash_mla_inputs( - M: int, # noqa: N803 - past_kv: int, - *, - K_per_query: int, # noqa: N803 - batch_size: int, - n_local_heads: int, - device: str, -) -> tuple[torch.Tensor, ...]: - """Build fp8 paged K cache + Q + indices + scheduler metadata. - - Layout = MODEL1_FP8Sparse (DSV4 NSA): d_qk=512 with 584-byte fp8 cache. - """ - M_per_req = M // batch_size if batch_size > 1 else M # noqa: N806 - # Per-request sequence length: each of the ``batch_size`` requests has - # ``M_per_req`` new query tokens over ``past_kv`` prefix (NOT the flattened - # total M), so cache_seqlens / indices mirror the owning module - # (prefix, isl, bs) row 1:1 instead of one giant bs*isl request. - full_s = M_per_req + past_kv - K_per_query = max(min(K_per_query, full_s), 1) # noqa: N806 - - # Q: (batch, M_per_req, n_local_heads, FMLA_D_QK=512) bf16 - q = torch.randn(batch_size, M_per_req, n_local_heads, FMLA_D_QK, dtype=torch.bfloat16, device=device) - - # K cache: SHARED across batch entries to avoid b-fold blowup. - blocks_per_req = (full_s + _dsv4_kv_page_size() - 1) // _dsv4_kv_page_size() - k_bf16 = torch.randn(blocks_per_req, _dsv4_kv_page_size(), 1, FMLA_D_QK, dtype=torch.bfloat16, device=device) - k_cache = _quantize_k_cache_model1(k_bf16) - - block_table = _expand_block_table(blocks_per_req, batch_size, device) - - # indices_in_kvcache: (batch, M_per_req, K_per_query) int32 — first K_per_query positions per Q - indices = _expand_indices(K_per_query, batch_size, M_per_req, device) - - # Pad indices to multiple of 64 (FlashMLA assertion) - if K_per_query % 64 != 0: - pad = 64 - K_per_query % 64 - pad_t = torch.full((batch_size, M_per_req, pad), -1, dtype=torch.int32, device=device) - indices = torch.cat([indices, pad_t], dim=-1).contiguous() - - cache_seqlens = torch.full((batch_size,), full_s, dtype=torch.int32, device=device) - - return q, k_cache, block_table, indices, cache_seqlens - - -# ═══════════════════════════════════════════════════════════════════════ -# Kernel 2 (HCA): flash_mla_with_kvcache -# ═══════════════════════════════════════════════════════════════════════ - - -def _bench_flash_mla_sparse( - M: int, # noqa: N803 - past_kv: int, - *, - K_per_query: int, # noqa: N803 - native_heads: int, - sliding_window: int, - batch_size: int = 1, - tp_size: int = 1, - device: str = "cuda:0", -) -> float: - """Benchmark sparse FlashMLA matching V4's V4 backend call shape. - - Real V4 HCA backend (deepseek_v4_backend_radix.py:1087+) passes: - - main K cache + indices = SWA window (128 fixed positions) - - extra K cache + indices = c128 (HCA) positions - Total K attended per query = SWA_WINDOW + extra_K_per_query. - - TP zero-pad (mirrors ``sglang/srt/models/deepseek_v4.py:847``): - 1. Projection produces ``q_local`` of shape (..., native_heads//tp, d_qk) — - the rank's actual computed heads. - 2. ``q_padded`` is allocated full (..., native_heads, d_qk) and the rank's - ``tp_slice`` is filled from ``q_local``; other heads are zeros. - 3. FlashMLA always receives the full native head count. - """ - from flash_mla import flash_mla_with_kvcache, get_mla_metadata - - # rank-local head count (what the upstream projection actually produces) - n_local_heads = max(1, native_heads // tp_size) - M_per_req = M // batch_size if batch_size > 1 else M # noqa: N806 - _full_s = M_per_req + past_kv # per-request sequence length (bs requests, M_per_req new tokens each) - - # Build main K cache + ``q_local`` (per-rank Q at n_local_heads) - q_local, k_cache_main, _, _, cache_seqlens = _build_flash_mla_inputs( - M, - past_kv, - K_per_query=K_per_query, - batch_size=batch_size, - n_local_heads=n_local_heads, - device=device, - ) - - # Zero-pad ``q_local`` to the full native head count before the FlashMLA - # call. Passing the unpadded TP-local head count trips ``Unsupported h_q``. - if n_local_heads == native_heads: - q = q_local - else: - q = torch.zeros(batch_size, M_per_req, native_heads, FMLA_D_QK, dtype=torch.bfloat16, device=device) - q[:, :, :n_local_heads, :] = q_local # rank-0's tp_slice - - # Kernel always sees full h_q. - n_local_heads = native_heads - swa_window = sliding_window - swa_indices = _expand_indices(swa_window, batch_size, M_per_req, device) - swa_topk_lengths = torch.full((batch_size,), swa_window, dtype=torch.int32, device=device) - - # Build extra K cache (c128 or c4) + extra indices - _pg = _dsv4_kv_page_size() - extra_blocks = (K_per_query + _pg - 1) // _pg - extra_k_bf16 = torch.randn(max(1, extra_blocks), _pg, 1, FMLA_D_QK, dtype=torch.bfloat16, device=device) - extra_k_cache = _quantize_k_cache_model1(extra_k_bf16) - extra_K = max(64, ((K_per_query + 63) // 64) * 64) # noqa: N806 - extra_indices = _expand_indices(extra_K, batch_size, M_per_req, device) - extra_topk_lengths = torch.full((batch_size,), K_per_query, dtype=torch.int32, device=device) - - sched_meta, _ = get_mla_metadata( - cache_seqlens=cache_seqlens, - num_q_tokens_per_head_k=M_per_req * n_local_heads, - num_heads_k=1, - num_heads_q=n_local_heads, - is_fp8_kvcache=True, - topk=swa_indices.size(-1), - ) - - softmax_scale = 1.0 / (FMLA_D_QK**0.5) - attn_sink = torch.zeros(n_local_heads, dtype=torch.float32, device=device) - - def kernel_fn(): - flash_mla_with_kvcache( - q=q, - k_cache=k_cache_main, - block_table=None, - cache_seqlens=None, - head_dim_v=V_HEAD_DIM, - tile_scheduler_metadata=sched_meta, - num_splits=None, - softmax_scale=softmax_scale, - causal=False, - is_fp8_kvcache=True, - indices=swa_indices, - topk_length=swa_topk_lengths, - attn_sink=attn_sink, - extra_k_cache=extra_k_cache, - extra_indices_in_kvcache=extra_indices, - extra_topk_length=extra_topk_lengths, - ) - - return _bench_cuda_graph(kernel_fn, allow_graph_fail=False, device=device) - - -# HCA and CSA are the SAME FlashMLA kernel (``_bench_flash_mla_sparse``) — they -# differ only in how many KV positions each query attends to: -# - hca_attn: all c128 positions -> full_s // 128 -# - csa_attn: the topk-selected c4 positions -> min(index_topk, full_s // 4) -# (csa uses the same sequential-index approximation as hca; the real topk -# scatter pattern is not reproducible in a kernel-level bench.) -_FMLA_K_PER_QUERY = { - "hca_attn": lambda full_s, sc: max(1, full_s // 128), - "csa_attn": lambda full_s, sc: max(1, min(sc.index_topk, full_s // 4)), -} - - # ═══════════════════════════════════════════════════════════════════════ # CSV write helper # ═══════════════════════════════════════════════════════════════════════ @@ -678,7 +418,7 @@ def _write_row( ) -> None: os.makedirs(os.path.dirname(os.path.abspath(perf_filename)) or ".", exist_ok=True) - mla_dtype = "bfloat16" if kernel in ("hca_attn", "csa_attn") else "fp8_e4m3" + mla_dtype = "fp8_e4m3" kv_cache_dtype = "fp8_e4m3" gemm_type = "fp8_block" @@ -745,29 +485,13 @@ def _guarded_bench(bench_fn: Callable[[], object], label: str): def _bench_sparse_kernel_shape(kernel, prefix, isl, bs, sc, device, *, topk_variant=None): """Bench one module ``(prefix, isl, bs)`` shape; return a list of ``(score_mode, latency_ms)`` rows. TP-independent (paged_mqa maps tokens to - b=M while preserving request-local causal lengths; FMLA pads to full native - heads), so this is computed once per shape. + b=M while preserving request-local causal lengths), so this is computed + once per shape. Single-latency kernels return one ``(None, latency)``; topk returns the phase-qualified flat-vs-representative pair whose matching-variant DELTA perf_database consumes (flat.latency - top_last.latency).""" M = bs * isl # noqa: N806 - if kernel in _FMLA_K_PER_QUERY: - # FMLA mirrors the module row 1:1: bs requests, each isl new query tokens - # over `prefix` past_kv. Pass batch_size=bs (so M_per_req=isl) and base - # K_per_query on the PER-REQUEST sequence length (isl+prefix), not the - # flattened total (bs*isl+prefix). - full_s = isl + prefix - lat = _bench_flash_mla_sparse( - M, - prefix, - K_per_query=_FMLA_K_PER_QUERY[kernel](full_s, sc), - native_heads=sc.num_attention_heads, - sliding_window=sc.sliding_window, - batch_size=bs, - device=device, - )["latency_ms"] - return [(None, lat)] if kernel == "topk": if topk_variant not in ("v1", "v2"): raise ValueError(f"topk requires variant v1 or v2, got {topk_variant!r}") @@ -843,7 +567,12 @@ def _sglang_chunked_prefill_size(): if _CHUNKED_PREFILL is None: try: from sglang.srt.model_executor.cuda_graph_config import default_cuda_graph_config - from sglang.srt.server_args import ServerArgs, get_device_memory_capacity + from sglang.srt.server_args import ServerArgs + + try: # 0.5.21: moved to utils/common.py:904 (server_args no longer re-exports it) + from sglang.srt.utils.common import get_device_memory_capacity + except ImportError: + from sglang.srt.server_args import get_device_memory_capacity except ModuleNotFoundError: from srt.model_executor.cuda_graph_config import default_cuda_graph_config from srt.server_args import ServerArgs, get_device_memory_capacity @@ -852,16 +581,23 @@ def _sglang_chunked_prefill_size(): gpu_mem = get_device_memory_capacity("cuda") except Exception: pass - sa = ServerArgs.__new__(ServerArgs) - sa.chunked_prefill_size = None - sa.cuda_graph_config = default_cuda_graph_config() - sa.tp_size = 1 - sa.device = "cuda" + chunked = None try: - sa._handle_gpu_memory_settings(gpu_mem) - except Exception: - pass # chunked_prefill_size is set first, before any model-dependent step - chunked = sa.chunked_prefill_size + sa = ServerArgs.__new__(ServerArgs) + except TypeError: + # sglang>=0.5.20: ServerArgs is a msgspec Struct (no bare __new__); the + # GPU-memory tiering below is sglang's own and needs no instance. + sa = None + if sa is not None: + sa.chunked_prefill_size = None + sa.cuda_graph_config = default_cuda_graph_config() + sa.tp_size = 1 + sa.device = "cuda" + try: + sa._handle_gpu_memory_settings(gpu_mem) + except Exception: + pass # chunked_prefill_size is set first, before any model-dependent step + chunked = sa.chunked_prefill_size if chunked is None and gpu_mem is not None: # Newer sglang (0.0.0.dev / >=0.5.x) refactored cuda_graph_max_bs/_bs # into a cuda_graph_config object that _handle_gpu_memory_settings @@ -887,7 +623,7 @@ def _sglang_chunked_prefill_size(): def _dsv4_context_derived_shapes(model_path): """DSV4 sparse-kernel context shapes derived from the csa/hca context MODULE INPUT sweep (same source as GLM5's dsa derive) — NOT read back from the - module CSV. paged_mqa_logits / topk / csa_attn / hca_attn therefore cover + module CSV. paged_mqa_logits / topk therefore cover the full-chunk isl the module drops at runtime; CSA and HCA share one grid (the module input validity does not depend on attn_kind). @@ -988,9 +724,9 @@ def run_dsv4_sparse_kernel_worker( """Single worker for the whole DSV4 sparse-op family (invoked by collect.py, one case per model from get_func). - Reads the kernel's owning module CSV(s) — paged_mqa_logits/topk/csa_attn ← - CSA, hca_attn ← HCA — dedups to the unique ``(prefix, isl, bs)`` module shapes - (the kernels are TP-independent), and benches each once. + Derives the kernel's shapes from its owning CSA module INPUT sweep + (paged_mqa_logits / topk ← CSA) — dedups to the unique ``(prefix, isl, bs)`` + module shapes (the kernels are TP-independent), and benches each once. ``_bench_sparse_kernel_shape`` returns one or more ``(score_mode, latency)`` rows per shape — one for the single-latency kernels, two for topk (variant-qualified flat + top_last) — each written as one row (tp_size=1).""" @@ -1131,11 +867,30 @@ def _bench_topk_512( uses planned v2. Both execute inside production CUDA graphs, so capture is mandatory here too. """ - from sglang.jit_kernel.dsv4.topk import ( - plan_topk_v2, - topk_transform_512, - topk_transform_512_v2, - ) + try: + from sglang.jit_kernel.dsv4.topk import ( + plan_topk_v2, + topk_transform_512, + topk_transform_512_v2, + ) + except ImportError: + # sglang>=0.5.21 moved the JIT ops to sglang.kernels.ops.attention.dsv4.topk and + # renamed the entry points: topk_transform_512 -> topk_transform_paged (CUDA branch = + # the same _jit_topk_v1_module().topk_transform, topk.py:104-125; the ROCm branch is the + # deepseek_v4_topk_transform_512 AOT op) and topk_transform_512_v2 -> + # topk_transform_paged_v2 (topk.py:222, identical positional signature, out_raw_indices + # now optional). Serving DSV4 dispatches exactly these two from + # indexer.topk_transform_paged_from_metadata (srt/layers/attention/dsv4/indexer.py:446-486, + # v2 iff metadata.use_topk_v2), so the v1/v2 calib lanes keep their meaning. + from sglang.kernels.ops.attention.dsv4.topk import ( + plan_topk_v2, + ) + from sglang.kernels.ops.attention.dsv4.topk import ( + topk_transform_paged as topk_transform_512, + ) + from sglang.kernels.ops.attention.dsv4.topk import ( + topk_transform_paged_v2 as topk_transform_512_v2, + ) if variant not in ("v1", "v2"): raise ValueError(f"unknown topk variant: {variant}") diff --git a/python/aisimulate/collector/sglang/dsa_prefill_graph.py b/python/aisimulate/collector/sglang/dsa_prefill_graph.py index fbbf68d724..425f831af0 100644 --- a/python/aisimulate/collector/sglang/dsa_prefill_graph.py +++ b/python/aisimulate/collector/sglang/dsa_prefill_graph.py @@ -9,6 +9,15 @@ model_executor/runner/prefill_cuda_graph_runner.py:111-318,370-427,765-941. This adapter delegates compilation, capture, padding, metadata and replay to that runtime. It excludes batch preparation and the model's other layers. + +FIXME(port): the adapter (dsa_prefill_graph) still follows the 0.5.14 +tc_piecewise runner. At 0.5.21 the collector keeps prefill graphs off +(collect_mla_module.load_model_runner), so graph_token_bucket returns None and +the adapter is not entered; any other prefill backend fails closed there. A +0.5.21 port needs the layer_boundary AttentionInputs import, the resolved +config through get_exec().graph (override, not a server_args copy), +mha_companion_layers on the runner and both context managers, and +support_pp=False; serving's default for DSA is the "breakable" backend. """ import copy @@ -35,12 +44,20 @@ def _release_owned_compile_hooks(hooks, previous_keys, module): def graph_token_bucket(model_runner, tokens: int) -> int | None: - config = model_runner.server_args.cuda_graph_config.prefill + """The prefill-graph bucket for ``tokens``, or None when the runner's prefill graph is off. + + 0.5.21 runners read the resolved config through get_exec() + (prefill_cuda_graph_runner.py:310-326); ServerArgs.cuda_graph_config stays None. + """ + from sglang.srt.runtime_context import get_exec + + config = get_exec().graph.cuda_graph_config.prefill if config.backend == "disabled": return None - if config.backend != "tc_piecewise": - raise ValueError(f"Unsupported SGLang module prefill graph backend: {config.backend}") - return next((size for size in sorted(config.bs or []) if size >= tokens), None) + raise ValueError( + f"SGLang {config.backend} prefill graph: the attention-only adapter follows the 0.5.14 " + "tc_piecewise runner and is not ported to 0.5.21 (see the module docstring)" + ) @contextmanager diff --git a/python/aisimulate/collector/sglang/glm5_dsa_sparse_modules.py b/python/aisimulate/collector/sglang/glm5_dsa_sparse_modules.py index 0e87f785c8..d266fbd60c 100644 --- a/python/aisimulate/collector/sglang/glm5_dsa_sparse_modules.py +++ b/python/aisimulate/collector/sglang/glm5_dsa_sparse_modules.py @@ -27,7 +27,10 @@ from __future__ import annotations -__compat__ = "sglang==0.5.14" +# 0.5.21 since 2026-10-01 (H20/sm90 collector port: op_smoke + path gates in the v0.5.21 image; +# findings hopper_sglang_collector_port_0514_to_0521_2026_10_01). +# One release per file (owner policy 2026-09-20); earlier code is `git log -- ` away. +__compat__ = "sglang==0.5.21" import os import sys diff --git a/python/aisimulate/collector/sglang/registry.py b/python/aisimulate/collector/sglang/registry.py index 086641ac4a..fd2409f2a1 100644 --- a/python/aisimulate/collector/sglang/registry.py +++ b/python/aisimulate/collector/sglang/registry.py @@ -32,6 +32,9 @@ get_func="get_computescale_test_cases", run_func="run_computescale", perf_filename=PerfFile.COMPUTESCALE, + # one measurement, two tables: the static quant row goes to + # scale_matrix_perf; finalize binds both to this producer + extra_perf_filenames=(PerfFile.SCALE_MATRIX,), ), OpEntry( op="mla_context", @@ -123,16 +126,21 @@ ), # GLM-5.2 skip-indexer layers (index_topk_freq>1): same shapes as the full # DSA module, but the per-layer indexer (mqa+topk+index-K store) is patched - # out so the captured cost is the reuse-layer cost. run_func derives a - # skip_indexer bool from the worker_perf_filename selector and passes it to - # the benchmark subprocess as an explicit arg (no env var). + # out so the captured cost is the reuse-layer cost. The rows go into the + # SAME dsa_*_module table as the full op, tagged by an op_name + # "_skip_indexer" suffix (that is what the SDK reads), so the entry declares + # that shared table and selects the mode with run_kwargs — the executor + # binds skip_indexer=True to run_func. (Until 2026-10-06 the entry declared + # a dsa_*_skip_indexer_perf.txt table nobody wrote and the worker sniffed + # the mode from that filename; --resume finalization trusted the + # declaration and failed every skip_indexer shard: job 469988017.) OpEntry( op="dsa_context_module_skip_indexer", module="collector.sglang.collect_mla_module", get_func="get_dsa_context_module_skip_indexer_test_cases", run_func="run_mla_module_worker", perf_filename=PerfFile.DSA_CONTEXT_MODULE, - worker_perf_filename=PerfFile.DSA_CONTEXT_MODULE_SKIP_INDEXER, + run_kwargs=(("skip_indexer", True),), # Hardware-validated on SM100, on SM103 via the B300 probe # (2026-07-13, pipeline 57747474: 8,506 rows, skip_indexer # trtllm+flashmla buckets clean, 1 error), and on SM90 via the @@ -151,7 +159,7 @@ get_func="get_dsa_generation_module_skip_indexer_test_cases", run_func="run_mla_module_worker", perf_filename=PerfFile.DSA_GENERATION_MODULE, - worker_perf_filename=PerfFile.DSA_GENERATION_MODULE_SKIP_INDEXER, + run_kwargs=(("skip_indexer", True),), unverified_sms=(120,), ), # DeepSeek-V4 module-level data (csa/hca x ctx/gen = 4 ops, 1 file each). @@ -225,28 +233,6 @@ # hardware-validate through the sanctioned plan. unverified_sms=(120,), ), - # The standalone HCA/CSA FMLA sub-kernel collectors require the - # ``flash_mla`` package, which the pinned 0.5.14 image does not ship (the - # serving path uses the bundled sgl-kernel FlashMLA instead); no platform - # has ever produced rows for them. Keep them parked as unverified until a - # wired backend exists, instead of letting the getter silently enumerate - # zero cases. - OpEntry( - op="dsv4_hca_attn_module", - module="collector.sglang.deepseekv4_sparse_modules", - get_func="get_dsv4_hca_attn_test_cases", - run_func="run_dsv4_sparse_kernel_worker", - perf_filename=PerfFile.DSV4_HCA_ATTN_MODULE, - unverified=True, - ), - OpEntry( - op="dsv4_csa_attn_module", - module="collector.sglang.deepseekv4_sparse_modules", - get_func="get_dsv4_csa_attn_test_cases", - run_func="run_dsv4_sparse_kernel_worker", - perf_filename=PerfFile.DSV4_CSA_ATTN_MODULE, - unverified=True, - ), # CSA topk_512 DELTA calibration (flat vs top_last) — feeds the # degenerate->representative topK correction in perf_database. OpEntry( @@ -300,10 +286,11 @@ perf_filename=PerfFile.GLM5_DSA_ATTN_MODULE, unverified_sms=(120,), ), - # MiniMax-M3 MSA sparse-attention modules — requires the msa-family image - # pin (framework_manifest sglang families.msa → official v0.5.16, the - # first release with models/minimax_m3.py; the module declares - # __compat__ = "sglang>=0.5.16"). Hardware-validated per SM: + # MiniMax-M3 MSA sparse-attention modules — need models/minimax_m3.py, + # first shipped in the official v0.5.16 release (the module declares + # __compat__ = "sglang>=0.5.16"; the 0.5.14-era framework_manifest carried + # an msa family pin to v0.5.16, retired when the default moved to 0.5.21 + # on 2026-10-01). Hardware-validated per SM: # SM90 (h100/h200 — SGLang's own M3 server-args override selects # fa3 + page 128 and the Triton sparse path, # arg_groups/overrides.py:521-537@v0.5.16), SM100/103 diff --git a/python/aisimulate/collector/sglang/runtime_compat.py b/python/aisimulate/collector/sglang/runtime_compat.py new file mode 100644 index 0000000000..622369274b --- /dev/null +++ b/python/aisimulate/collector/sglang/runtime_compat.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Process-wide sglang runtime config for the module collectors, across sglang generations. + +sglang 0.5.21 (runtime_context "bags"): a worker process must ``publish(server_args, role=..., +ranks=SpawnRanks(...))`` first and the config initializers then take NO arguments +(``benchmark/one_batch.py:893-902`` @0.5.21: publish -> initialize_moe_config() -> +initialize_fp8_gemm_config() -> initialize_fp4_gemm_config()); ``ModelRunner.__init__`` no longer +accepts the parallel geometry (``model_executor/model_runner.py``: model_config, mem_fraction_static, +gpu_id, nccl_port, server_args, ...) — it reads ``get_parallel()``. +sglang 0.5.16: the initializers take ``server_args`` and the geometry travels as +``ps=ParallelState.trivial(gpu_id)`` (parallel_state_wrapper.py:6-24). +Older pins: initializers take ``server_args`` and ModelRunner takes tp_rank/tp_size/... keywords. + +This only changes HOW the same engine state is published (layer_permissions.md API-compat shim +rule); the kernels the runner selects are untouched. +""" + +from __future__ import annotations + +import inspect + + +def init_runtime_config( + server_args, + gpu_id: int, + *, + nccl_port: int | None = None, + model_config=None, + tp_rank: int = 0, + tp_size: int | None = None, +) -> dict: + """Run the framework's own config initializers and return the ModelRunner parallel kwargs. + + 0.5.21 additionally needs the distributed world built BEFORE ModelRunner (model_runner.py + init_torch_distributed -> bootstrap.measure_pre_model_load_memory reads get_world_group()): + scheduler.py:525-530 and benchmark/one_batch.py load_model run bootstrap.init_parallel_runtime + (publish -> groups) then bootstrap.init_layer_runtime(model_config). Pass nccl_port + model_config. + """ + from sglang.srt.layers.moe import initialize_moe_config + from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config + from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config + + if inspect.signature(initialize_moe_config).parameters: # <= 0.5.20: seeded from ServerArgs directly + initialize_moe_config(server_args) + initialize_fp8_gemm_config(server_args) + initialize_fp4_gemm_config(server_args) + try: + from sglang.srt.distributed.parallel_state_wrapper import ParallelState + + return {"ps": ParallelState.trivial(gpu_id=gpu_id)} + except ImportError: + return { + "tp_rank": tp_rank, + "tp_size": tp_size or getattr(server_args, "tp_size", 1), + "pp_rank": 0, + "pp_size": 1, + "moe_ep_rank": 0, + "moe_ep_size": 1, + } + from sglang.srt.runtime_context import SpawnRanks, publish, spawn_world_rank + + publish( + server_args, + role="scheduler", + ranks=SpawnRanks(world_rank=spawn_world_rank(server_args, tp_rank=tp_rank, pp_rank=0), gpu_id=gpu_id), + ) + initialize_moe_config() + initialize_fp8_gemm_config() + initialize_fp4_gemm_config() + from sglang.srt.distributed import bootstrap + + if nccl_port is None or model_config is None: + raise RuntimeError( + "sglang>=0.5.21 runner init needs nccl_port and model_config " + "(bootstrap.init_parallel_runtime/init_layer_runtime)" + ) + try: + from sglang.srt.runtime_context import get_device + + device = get_device().device + except (ImportError, AttributeError): + device = getattr(server_args, "device", "cuda") + bootstrap.init_parallel_runtime(server_args=server_args, device=device, dist_port=nccl_port) + bootstrap.init_layer_runtime(model_config=model_config) + return {} + + +def free_local_port() -> int: + """A port the OS reports free right now: an arithmetic guess can collide between + concurrent per-GPU workers, and the collision surfaces as a framework failure.""" + import socket + + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +_OFFLINE_PUBLISHED = False + + +def ensure_offline_runtime_published(**server_args_fields) -> bool: + """Publish a minimal process-wide runtime config for the kernel-level collectors + that build attention backends on a mock runner (no engine, no model). + + sglang>=0.5.20 projects ``get_exec()`` / ``get_parallel()`` from a published + ServerArgs; the backends read them (flashattention_backend.py:320,349 get_exec(). + deterministic; parallel_state world_group). sglang's own unit tests publish an + offline record exactly this way (``ServerArgs(model_path="dummy", **fields)`` + + ``publish(role=..., ranks=SpawnRanks(world_rank=0))``, test/test_utils.py:1990-2023 + @0.5.21; the resolution pipeline returns early on the dummy/absent-model path). + Older sglang has no ``publish`` -> no-op. Returns True when a publish happened. + """ + global _OFFLINE_PUBLISHED + if _OFFLINE_PUBLISHED: + return True + try: + from sglang.srt.runtime_context import SpawnRanks, publish + from sglang.srt.server_args import ServerArgs + except ImportError: + return False + publish(ServerArgs(model_path="dummy", **server_args_fields), role="scheduler", ranks=SpawnRanks(world_rank=0)) + _OFFLINE_PUBLISHED = True + return True + + +def attach_kv_index_translator(mock_runner) -> None: + """sglang>=0.5.21 attention backends take ``model_runner.kv_index_translator`` + (flashattention_backend.py:205; built by ModelRunner.init_kv_index_translator at + model_runner.py:860-868 from req_to_token / allocator / pool / page_size / device). + Build the real one when the mock runner carries those objects; otherwise a + non-translating stand-in (the backends only consult ``is_translating`` on the + non-DCP, non-hybrid paths these kernel collectors exercise). No-op on older sglang. + """ + try: + from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator + except ImportError: + return + try: + mock_runner.kv_index_translator = KVIndexTranslator( + req_to_token=mock_runner.req_to_token_pool.req_to_token, + token_to_kv_pool_allocator=mock_runner.token_to_kv_pool_allocator, + token_to_kv_pool=mock_runner.token_to_kv_pool, + page_size=getattr(mock_runner, "page_size", None) or getattr(mock_runner.server_args, "page_size", 1) or 1, + device=str(mock_runner.device), + ) + except Exception as exc: + from types import SimpleNamespace + + print( + f"[sglang-compat] KVIndexTranslator not constructible on the mock runner ({type(exc).__name__}: {exc}); " + "using a non-translating stand-in" + ) + mock_runner.kv_index_translator = SimpleNamespace(is_translating=False, reads_are_translated=False) + + +def resolved_arg(server_args, name: str, default=None): + """A ServerArgs field as serving resolves it: the raw record first, else the + sglang>=0.5.21 resolving view (arg_groups/model_override_base.py:120) where derived + values (page_size, chunked_prefill_size, cuda_graph_config, ...) live.""" + value = getattr(server_args, name, None) + if value is not None: + return value + try: + from sglang.srt.arg_groups.model_override_base import resolving_view + + value = getattr(resolving_view(server_args), name, None) + except ImportError: + value = None + return default if value is None else value + + +def reset_expert_location_metadata() -> None: + """Return the publish-once EPLB expert-location metadata to its pre-init state, so a + second ModelRunner in the same worker can initialise. sglang 0.5.21 keeps it in + ``get_resources().expert_location_metadata`` (runtime_context.py:621) and + ``set_global_expert_location_metadata`` asserts it is None (eplb/expert_location.py: + 484-489); older pins kept the ``_global_expert_location_metadata`` module global.""" + try: + from sglang.srt.runtime_context import get_resources + except ImportError: + from sglang.srt.eplb import expert_location + + expert_location._global_expert_location_metadata = None + return + get_resources().expert_location_metadata = None + + +def causal_conv1d_uses_int64_offsets() -> bool: + """Whether the installed causal_conv1d Triton kernel computes token offsets in + int64. The 0.5.14 kernel used int32 pointer arithmetic (causal_conv1d_triton.py:373-379, + collector guards at total_tokens*conv_channels >= 2**31); sglang>=0.5.21 casts the + sequence/token offsets to tl.int64 (kernels/ops/mamba/causal_conv1d_triton.py:95,185), + so that limit no longer exists there. Source-probed, never version-sniffed.""" + import importlib + import inspect + + for name in ( + "sglang.kernels.ops.mamba.causal_conv1d_triton", + "sglang.srt.layers.attention.mamba.causal_conv1d_triton", + ): + try: + src = inspect.getsource(importlib.import_module(name)) + except Exception: + continue + return "tl.int64" in src and "stride_x_token" in src + return False diff --git a/python/aisimulate/collector/sglang/runtime_limits.py b/python/aisimulate/collector/sglang/runtime_limits.py index a9caa5ed80..cfd7445bc6 100644 --- a/python/aisimulate/collector/sglang/runtime_limits.py +++ b/python/aisimulate/collector/sglang/runtime_limits.py @@ -149,6 +149,23 @@ def sglang_dsa_mqa_logits_chunking_supported() -> bool: has_chunk_kernel = "logits_chunk" in source if has_chunk_decision and has_chunk_loop and has_chunk_kernel: return True + # sglang>=0.5.21 restructured the contract: the free-memory row chunking moved + # out of Indexer into dsa_indexer_kpool._mqa_logits_row_chunks (:85, driven by + # mqa_logits_should_chunk / mqa_logits_rows_per_chunk / mqa_logits_budget_bytes) + # and the SM100 batch>num_sms split lives in Indexer._get_topk_ragged's local + # _chunked_fp8_paged_mqa_logits (dsa_indexer.py:1027). Same behaviour the + # collector relies on (logits never exceed the free-memory budget), new shape. + try: + kpool = importlib.import_module("sglang.srt.layers.attention.dsa.dsa_indexer_kpool") + indexer_mod = importlib.import_module("sglang.srt.layers.attention.dsa.dsa_indexer") + indexer_cls = getattr(indexer_mod, "Indexer", None) + row_chunks = getattr(kpool, "_mqa_logits_row_chunks", None) + if indexer_cls is not None and row_chunks is not None and getattr(indexer_cls, "_get_topk_ragged", None): + source = inspect.getsource(row_chunks) + if "mqa_logits_should_chunk" in source and "mqa_logits_rows_per_chunk" in source: + return True + except Exception: + pass raise RuntimeError("SGLang DSA MQA-logits chunking source contract was not detected") @@ -156,11 +173,14 @@ def dsa_indexer_prefill_shape_is_supported(batch_size: int, seq_len: int) -> boo """Validate dimensions without inferring the native forward mode. A one-token EXTEND remains prefill. SGLang 0.5.14's native TC runner - bucket-pads it (49e384ce, prefill_cuda_graph_runner.py:480-524), and its - DSA backend selects the sparse path in that context. The old eager-only - failure does not establish a native prefill graph exclusion. + bucket-pads it (49e384ce, prefill_cuda_graph_runner.py:480-524). At 0.5.21 + the collector times DSA prefill eagerly (prefill graphs off, as serving runs + a one-token extend too: bucket 4 is more than twice its size, + prefill_cuda_graph_runner.py:157,1286-1290), and the eager DSA prefill + indexer illegal-accessed on that shape. FIXME(re-verify): admit seq_len=1 + once a 0.5.21 run shows the eager one-token extend no longer faults. """ - return batch_size > 0 and seq_len > 0 + return batch_size > 0 and seq_len > 1 def dsa_indexer_total_kv_tokens_supported( @@ -183,6 +203,16 @@ def required_prefill_extend_tokens(batch_size: int, seq_len: int) -> int: def runtime_chunk_size(model_runner) -> int: server_args = getattr(model_runner, "server_args", None) sglang_chunk = getattr(server_args, "chunked_prefill_size", None) if server_args else None + if not (isinstance(sglang_chunk, int) and sglang_chunk > 0) and server_args is not None: + # sglang>=0.5.21: derived values are read through the resolving view + # (arg_groups/model_override_base.py:120 resolving_view; one_batch.load_model + # reads cfg.mem_fraction_static the same way), not off the raw record. + try: + from sglang.srt.arg_groups.model_override_base import resolving_view + + sglang_chunk = getattr(resolving_view(server_args), "chunked_prefill_size", None) + except ImportError: + pass if isinstance(sglang_chunk, int) and sglang_chunk > 0: return sglang_chunk raise RuntimeError("SGLang did not initialize server_args.chunked_prefill_size") diff --git a/python/aisimulate/collector/trtllm/collect_attn.py b/python/aisimulate/collector/trtllm/collect_attn.py index 8d154fb76f..9968d56ae7 100644 --- a/python/aisimulate/collector/trtllm/collect_attn.py +++ b/python/aisimulate/collector/trtllm/collect_attn.py @@ -27,13 +27,31 @@ ) from collector.helper import benchmark_with_power, get_sm_version, log_perf from collector.registry_types import PerfFile -from tensorrt_llm._torch.attention_backend import TrtllmAttentionMetadata -from tensorrt_llm._torch.attention_backend.interface import ( - AttentionRuntimeFeatures, - PositionalEmbeddingParams, - RopeParams, -) -from tensorrt_llm._torch.attention_backend.utils import create_attention + +# trtllm >=1.3.0rc29 moved _torch.attention_backend.* -> _torch.attention.backends.* and +# _torch.modules.fused_moe -> _torch.moe.fused_moe (the old package root is a deprecation shim +# without submodules). Path-only compat: same classes, same kernels (layer_permissions.md +# 'API-compat shims may only change HOW the same kernel is constructed'). +try: + from tensorrt_llm._torch.attention.backends import TrtllmAttentionMetadata +except ModuleNotFoundError: # < rc29 layout + from tensorrt_llm._torch.attention_backend import TrtllmAttentionMetadata +try: + from tensorrt_llm._torch.attention.backends.interface import ( + AttentionRuntimeFeatures, + PositionalEmbeddingParams, + RopeParams, + ) +except ModuleNotFoundError: # < rc29 layout + from tensorrt_llm._torch.attention_backend.interface import ( + AttentionRuntimeFeatures, + PositionalEmbeddingParams, + RopeParams, + ) +try: + from tensorrt_llm._torch.attention.backends.utils import create_attention +except ModuleNotFoundError: # < rc29 layout + from tensorrt_llm._torch.attention_backend.utils import create_attention from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.functional import PositionEmbeddingType @@ -41,7 +59,7 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig -__compat__ = "trtllm>=1.3.0rc20" +__compat__ = "trtllm==1.3.0rc29" def _skip_trtllm_sm120_fp8_context_fmha( @@ -82,6 +100,28 @@ def _skip_trtllm_sm120_fp8_context_fmha( ) +def _deployment_tokens_per_block() -> int: + """KV page the deployment actually serves with, not a collector constant. + + The generator renders every trtllm engine with ``kv_cache_config.tokens_per_block: 32`` + (TRT-LLM's own KvCacheConfig default); the collector used to pin 64 here, which on + Blackwell selects a different FMHA cubin family — the sm100 kernel names carry the page + (fmhaSm100a…PagedKvCausalP64… vs serving's …P32…), so the sm100 attention gates + DIVERGED on nothing but the page while the sm90 names hide it (B200, 2026-10-04, + opharness finding sm100_sglang_trtllm_gate_verdicts_2026_10_04). Ask the framework + for its default; ``AIS_TRTLLM_TOKENS_PER_BLOCK`` is the A/B hook. + """ + override = os.environ.get("AIS_TRTLLM_TOKENS_PER_BLOCK") + if override: + return int(override) + try: + from tensorrt_llm.llmapi import KvCacheConfig + + return int(KvCacheConfig().tokens_per_block) + except Exception: # older builds without the pydantic config + return 32 + + def run_attention_torch( batch_size, input_len, @@ -98,7 +138,12 @@ def run_attention_torch( device="cuda:0", ): device = torch.device(device) - torch.set_default_device(device) + # No torch.set_default_device(device): the FLASHINFER sub-backend builds its HOST block + # tables with torch.zeros(..., pin_memory=True) (attention/backends/flashinfer.py:1566 + # @1.3.0rc29); under a CUDA default device that allocation lands on the GPU and pinning + # raises "Only dense CPU tensors can be pinned" — 2,942 FLASHINFER decode cases (every + # Gemma-4 head_dim 256/512 x window 1024/0 cell) on b200 rc29, job 469988228. Serving + # never sets a default device; every tensor this collector creates names its device. torch.cuda.set_device(device) if attn_backend_name not in {"TRTLLM", "FLASHINFER"}: @@ -118,24 +163,67 @@ def run_attention_torch( # supported-arch set against fmhaRunner.cuh on the next framework version bump. if is_flashinfer and get_sm_version() not in (100, 103): raise ValueError( - f"FlashInfer trtllm-gen FMHA is Blackwell-only (SM100/103); Gemma4 dense " + f"FIXME(kernel-limit): FlashInfer trtllm-gen FMHA is Blackwell-only (SM100/103); Gemma4 dense " f"attention has no kernel on SM{get_sm_version()} " f"(fmhaRunner.cuh:37 + modeling_gemma4.py:270 @1.3.0rc20)" ) if is_flashinfer and use_fp8_context_fmha: # Mirror of the sglang collector's compute-dtype labeling rule: - # TRT-LLM 1.3.0rc20 FlashInferAttention has no FP8 FMHA compute path, + # TRT-LLM FlashInferAttention has no FP8 FMHA compute path, # including under the serving-pinned trtllm-gen sub-backend set below. # trtllm-gen's only fp8 handling is the fp8 KV cache: it casts Q to the # KV dtype to reuse the QkvE4m3OBfloat16 cubins but the FMHA output is - # still BF16 (attention_backend/flashinfer.py:1789-1791@1.3.0rc20). The + # still BF16 (attention/backends/flashinfer.py:2740-2746@1.3.0rc29; :1789-1791@1.3.0rc20). The # collector's fp8-context-FMHA case (quant_algo=FP8 + out_scale) is a # TRTLLM-backend contract flashinfer does not consume, so an fp8-labeled # flashinfer row would record BF16 compute under an fp8 label. Fail closed. - raise ValueError("TRT-LLM 1.3.0rc20 FlashInferAttention has no FP8 FMHA compute path") + raise ValueError( + "FIXME(kernel-limit): TRT-LLM FlashInferAttention has no FP8 FMHA compute path " + "(attention/backends/flashinfer.py:2740-2746 @1.3.0rc29)" + ) + # FIXME(kernel-limit): trtllm-gen has no FMHA for head_dim 192 on SM100/103 — context + # ("Missing TRTLLM-GEN kernel (context): headDimQk=192", 2,904 b200 rc29 cases) and decode + # ("Missing TRTLLM-GEN kernel (decode)", 1,624 b200 / 1,641 b300). Serving routes the same + # kernel. Classified before the layer is built; re-verify on the next TRT-LLM bump. + if head_dim == 192 and get_sm_version() in (100, 103): + raise ValueError( + "FIXME(kernel-limit): TRT-LLM trtllm-gen FMHA has no head_dim=192 kernel on SM100/103 " + f"({'context' if is_context_phase else 'decode'}; trtllm_fmha_kernel_launcher.cu:320)" + ) + # FIXME(kernel-limit): SM90 FP8-KV decode crashes ("CUDA error: unspecified launch failure" + # + SIGABRT) in two measured envelopes (h200 rc29 attention_generation shard, collector 3c6b058f, + # 2026-10-07: 512 launch failures + 169 illegal accesses, each followed by a worker SIGABRT; the + # shard's 55,047 passing rows bound them exactly): + # (a) num_heads 96 / kv_heads 4 (GQA 24): EVERY head_dim (64/128/256), every batch, isl + # and window — 512 cases, 0 passing rows; + # (b) head_dim 256 at GQA 32 (32/1, 64/2, 128/4), full attention (window 0): isl >= 255 + # — 169 cases; isl <= 127 passes at every batch. + # GQA 32/64 at other head_dims and GQA 24 never crashed and keep collecting (the first + # version of this guard, ed9109ec, refused every GQA >= 24 cell and lost 4,119 good rows: + # pipeline 72346027). Same family as the rc23 campaign's "decode fails for fp8 KV at GQA + # ratio 24/32" (finding campaign_1002 (3)). Re-verify on the next TRT-LLM bump. + if not is_context_phase and use_fp8_kv_cache and get_sm_version() == 90 and num_key_value_heads: + gqa_ratio = num_heads // num_key_value_heads + if (num_heads, num_key_value_heads) == (96, 4): + raise ValueError( + "FIXME(kernel-limit): TRT-LLM SM90 FP8-KV decode crashes for 96 q-heads over 4 kv-heads " + f"(head_dim={head_dim}); launch failure + SIGABRT on every cell" + ) + if head_dim == 256 and gqa_ratio == 32 and attention_window_size == 0 and input_len >= 255: + raise ValueError( + "FIXME(kernel-limit): TRT-LLM SM90 FP8-KV decode crashes for head_dim 256 at GQA ratio 32 " + f"with kv length >= 255 (num_heads={num_heads}, kv_heads={num_key_value_heads}, isl={input_len})" + ) # if XQA JIT is enabled, the context phase will also trigger XQA prepare which causes the error # with specifc q/kv head and seq setting. + # NOTE (1.3.0rc23, verified 2026-09-24): this env no longer exists in the + # binaries (only TRTLLM_FORCE_XQA / TRTLLM_XQA_BLOCKS_PER_SEQUENCE do). + # XQA vs MMHA for decode is the op's own heuristic ("JIT XQA is not + # used: maybe no performance gain" at short KV; XQA selected at long KV) — + # the collector follows it per (batch, kv_len) cell exactly as serving + # does, which path_diff confirmed at kv 4095 (XQA) vs kv 1 (MMHA). Kept + # for older builds that still read it. if is_context_phase: os.environ["TRTLLM_ENABLE_XQA_JIT"] = "0" else: @@ -144,7 +232,7 @@ def run_attention_torch( layer_idx = 0 world_size = 1 tp_size = 1 - tokens_per_block = 64 + tokens_per_block = _deployment_tokens_per_block() warming_up = 10 test_ite = 6 output_len = 1 @@ -352,7 +440,10 @@ def run_attention_torch( # with backend-agnostic kwargs (pyexecutor/model_engine.py:1784, # 1818-1830@1.3.0rc20); ``workspace`` is a TrtllmAttentionMetadata-only # field, flashinfer manages its own workspace_buffer. - from tensorrt_llm._torch.attention_backend.flashinfer import FlashInferAttentionMetadata + try: + from tensorrt_llm._torch.attention.backends.flashinfer import FlashInferAttentionMetadata + except ModuleNotFoundError: # < rc29 layout + from tensorrt_llm._torch.attention_backend.flashinfer import FlashInferAttentionMetadata attn_metadata = FlashInferAttentionMetadata(**metadata_kwargs) else: @@ -368,10 +459,10 @@ def run_attention_torch( else: num_tokens = batch_size - sinks = torch.randn(num_heads, dtype=torch.float32) if head_dim == 64 else None - q = torch.randn([num_tokens, num_heads * head_dim]).bfloat16().to(torch.device(device)) - k = torch.randn([num_tokens, num_key_value_heads * head_dim]).bfloat16().to(torch.device(device)) - v = torch.randn([num_tokens, num_key_value_heads * head_dim]).bfloat16().to(torch.device(device)) + sinks = torch.randn(num_heads, dtype=torch.float32, device=device) if head_dim == 64 else None + q = torch.randn([num_tokens, num_heads * head_dim], dtype=torch.bfloat16, device=device) + k = torch.randn([num_tokens, num_key_value_heads * head_dim], dtype=torch.bfloat16, device=device) + v = torch.randn([num_tokens, num_key_value_heads * head_dim], dtype=torch.bfloat16, device=device) if is_flashinfer: # Serving splits Q/K/V for backends without fused-QKV support # (modules/attention.py:619,641-646@1.3.0rc20; FlashInferAttention @@ -392,6 +483,20 @@ def run_attention_torch( "attention_sinks": sinks, "out_scale": out_scale, } + # FIXME(kernel-limit): H20 campaign 2026-10-02 (1.3.0rc29, TRTLLM backend, + # attention/backends/fmha/fallback.py:122): the context FMHA faults with + # cudaErrorIllegalAddress exactly when the LAST sequence's packed Q/O byte + # offset (batch-1) * input_len * num_heads * head_dim * 2 >= 2**31 — + # verified cell-by-cell on the (96,*,256), (64,*,256) and (96,8,128) grids + # (b16 s3072 h96 d256 crashes, b8 s6144 passes at the same token count), + # i.e. a 32-bit per-sequence offset inside the kernel. Serving reaches such + # batches only with max_num_tokens >= ~44k tokens per context forward. + # Unverified against the kernel source; cases above the boundary keep + # failing into the classified log until the claim is confirmed (or a bump + # fixes it). Generation has a second family: fp8 KV with GQA ratio 24 or + # 32 (96/4 at d64/128/256; 32/1, 64/2, 128/4 at d256) fails with + # cudaErrorLaunchFailure / IMA at every batch and KV length while the same + # shapes pass with a bf16 KV cache. attn.forward(*forward_args, attn_metadata, **forward_kwargs) # Use benchmark_with_power context manager diff --git a/python/aisimulate/collector/trtllm/collect_attn_encoder.py b/python/aisimulate/collector/trtllm/collect_attn_encoder.py index 7ba7f9a6de..ea029639ba 100644 --- a/python/aisimulate/collector/trtllm/collect_attn_encoder.py +++ b/python/aisimulate/collector/trtllm/collect_attn_encoder.py @@ -12,7 +12,7 @@ Quant: bf16 only. TRT-LLM upstream does not support fp8 on the encoder path. """ -__compat__ = "trtllm>=1.3.0rc20" +__compat__ = "trtllm==1.3.0rc29" import os @@ -21,12 +21,29 @@ from collector.case_generator import get_attention_encoder_head_configs, get_attention_encoder_shape_sweeps from collector.helper import benchmark_with_power, log_perf from collector.registry_types import PerfFile -from tensorrt_llm._torch.attention_backend import TrtllmAttentionMetadata -from tensorrt_llm._torch.attention_backend.interface import ( - AttentionRuntimeFeatures, - PredefinedAttentionMask, -) -from tensorrt_llm._torch.attention_backend.utils import create_attention + +# trtllm >=1.3.0rc29 moved _torch.attention_backend.* -> _torch.attention.backends.* and +# _torch.modules.fused_moe -> _torch.moe.fused_moe (the old package root is a deprecation shim +# without submodules). Path-only compat: same classes, same kernels (layer_permissions.md +# 'API-compat shims may only change HOW the same kernel is constructed'). +try: + from tensorrt_llm._torch.attention.backends import TrtllmAttentionMetadata +except ModuleNotFoundError: # < rc29 layout + from tensorrt_llm._torch.attention_backend import TrtllmAttentionMetadata +try: + from tensorrt_llm._torch.attention.backends.interface import ( + AttentionRuntimeFeatures, + PredefinedAttentionMask, + ) +except ModuleNotFoundError: # < rc29 layout + from tensorrt_llm._torch.attention_backend.interface import ( + AttentionRuntimeFeatures, + PredefinedAttentionMask, + ) +try: + from tensorrt_llm._torch.attention.backends.utils import create_attention +except ModuleNotFoundError: # < rc29 layout + from tensorrt_llm._torch.attention_backend.utils import create_attention from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig diff --git a/python/aisimulate/collector/trtllm/collect_computescale.py b/python/aisimulate/collector/trtllm/collect_computescale.py index a052b60a99..667de7d235 100644 --- a/python/aisimulate/collector/trtllm/collect_computescale.py +++ b/python/aisimulate/collector/trtllm/collect_computescale.py @@ -7,15 +7,22 @@ quantization for each YAML-backed shape. The reported latency isolates the scale-computation portion so support matrix data can track the cost separately from the static quantize kernel. + +One measurement, two tables — the registry entry declares the second one as +``extra_perf_filenames`` and the executor hands both names in, so finalize +binds both to this producer's checkpoint: + perf_filename -> computescale_perf.txt (dynamic quant minus static quant) + extra_perf_filenames[0] -> scale_matrix_perf.txt (the static per-tensor quant alone) """ -__compat__ = "trtllm>=1.3.0rc20" +__compat__ = "trtllm==1.3.0rc29" import tensorrt_llm import torch -from case_generator import get_compute_scale_case_specs +from collector.case_generator import get_compute_scale_case_specs +from collector.helper import benchmark_with_power, get_sm_version, log_perf -from helper import benchmark_with_power, get_sm_version, log_perf +_OUTSIDE_LOOP_COUNT = 5 # to reduce impact of L2 cache hit def get_computescale_test_cases(): @@ -30,90 +37,63 @@ def get_computescale_test_cases(): return test_cases -def run_computescale(m, k, *, perf_filename, device="cuda:0"): +def _setup(m, k, device): device = torch.device(device) torch.cuda.set_device(device) torch.set_default_device(device) + x = torch.randn((m, k), dtype=torch.bfloat16).to(device) + return device, x - dtype = torch.bfloat16 - x = torch.randn((m, k), dtype=dtype).to(torch.device(device)) - scale = torch.tensor([1.0], dtype=torch.float32, device=device) - outside_loop_count = 5 # to reduce impact of L2 cache hit - - # Build op lists for both dynamic and static quantization - dynamic_op_list = [] - static_op_list = [] - for _ in range(outside_loop_count): - dynamic_op_list.append(torch.ops.tensorrt_llm.quantize_e4m3_per_tensor) - static_op_list.append(torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor) - - # Benchmark dynamic quantization (compute scale + scale matrix) - def dynamic_kernel_func(): - for op in dynamic_op_list: - op(x) - - with benchmark_with_power( - device=device, - kernel_func=dynamic_kernel_func, - repeat_n=1, - ) as dynamic_results: +def _bench_dynamic(device, x): + # dynamic quantization = compute scale + scale matrix + def kernel_func(): + for _ in range(_OUTSIDE_LOOP_COUNT): + torch.ops.tensorrt_llm.quantize_e4m3_per_tensor(x) + + with benchmark_with_power(device=device, kernel_func=kernel_func, repeat_n=1) as results: pass + return results["latency_ms"] / _OUTSIDE_LOOP_COUNT, results["power_stats"] - dynamic_latency = dynamic_results["latency_ms"] / outside_loop_count - # Benchmark static quantization (scale matrix only) - def static_kernel_func(): - for op in static_op_list: - op(x, scale) +def _bench_static(device, x): + # static quantization = scale matrix only + scale = torch.tensor([1.0], dtype=torch.float32, device=device) + + def kernel_func(): + for _ in range(_OUTSIDE_LOOP_COUNT): + torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(x, scale) - with benchmark_with_power( - device=device, - kernel_func=static_kernel_func, - repeat_n=1, - ) as static_results: + with benchmark_with_power(device=device, kernel_func=kernel_func, repeat_n=1) as results: pass + return results["latency_ms"] / _OUTSIDE_LOOP_COUNT, results["power_stats"] - static_latency = static_results["latency_ms"] / outside_loop_count - # compute_scale latency = dynamic - static - compute_scale_latency = dynamic_latency - static_latency - compute_scale_latency = max(0.0, compute_scale_latency) +def run_computescale(m, k, *, perf_filename, extra_perf_filenames, device="cuda:0"): + (scale_matrix_filename,) = extra_perf_filenames + device, x = _setup(m, k, device) + dynamic_latency, dynamic_power = _bench_dynamic(device, x) + static_latency, static_power = _bench_static(device, x) + compute_scale_latency = max(0.0, dynamic_latency - static_latency) - # Log compute_scale performance log_perf( - item_list=[ - { - "m": m, - "k": k, - "quant_dtype": "fp8", - "latency": compute_scale_latency, - } - ], + item_list=[{"m": m, "k": k, "quant_dtype": "fp8", "latency": compute_scale_latency}], framework="TRTLLM", version=tensorrt_llm.__version__, device_name=torch.cuda.get_device_name(device), op_name="compute_scale", kernel_source="torch_ops", perf_filename=perf_filename, - power_stats=dynamic_results["power_stats"], + power_stats=dynamic_power, ) - # Log scale_matrix performance log_perf( - item_list=[ - { - "m": m, - "k": k, - "quant_dtype": "fp8", - "latency": static_latency, - } - ], + item_list=[{"m": m, "k": k, "quant_dtype": "fp8", "latency": static_latency}], framework="TRTLLM", version=tensorrt_llm.__version__, device_name=torch.cuda.get_device_name(device), op_name="scale_matrix", kernel_source="torch_ops", - perf_filename="scale_matrix_perf.txt", - power_stats=static_results["power_stats"], + perf_filename=scale_matrix_filename, + power_stats=static_power, ) diff --git a/python/aisimulate/collector/trtllm/collect_dsv4_attn.py b/python/aisimulate/collector/trtllm/collect_dsv4_attn.py index 95f8fd30cf..a1c9c500b5 100644 --- a/python/aisimulate/collector/trtllm/collect_dsv4_attn.py +++ b/python/aisimulate/collector/trtllm/collect_dsv4_attn.py @@ -41,7 +41,7 @@ from __future__ import annotations -__compat__ = "trtllm>=1.3.0rc21" +__compat__ = "trtllm==1.3.0rc29" import gc import json @@ -57,12 +57,12 @@ try: from registry_types import PerfFile - from helper import benchmark_with_power, log_perf + from helper import benchmark_with_power, get_sm_version, log_perf except ModuleNotFoundError: sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from registry_types import PerfFile - from helper import benchmark_with_power, log_perf + from helper import benchmark_with_power, get_sm_version, log_perf ARCHITECTURE = "DeepseekV4ForCausalLM" @@ -161,6 +161,31 @@ def _drop(reason: str) -> None: return shapes +# Perf-row / case label of each KV pool dtype = the SDK's KVCacheQuantMode member name. +_KV_CACHE_DTYPE_LABELS = {"bf16": "bfloat16", "bfloat16": "bfloat16", "fp8": "fp8"} + + +def _serving_kv_cache_dtype() -> str: + """KV pool dtype the deployment actually serves DSV4 with on this SM. + + Hopper (and the documented SM120/121 path) serve the packed fp8_ds_mla sparse-MLA pool. + On SM100 the generator's same facts (kv_cache_config.dtype fp8_ds_mla, page 128) are NOT + servable with 1.3.0rc29: the manager first demands tokens_per_block=256, then the fp8_ds_mla + sparse GENERATION step has no FMHA library ("No TRT-LLM attention FMHA library supports this + request"); the only configuration that boots is kv dtype auto -> bf16 (DeepseekV4CacheManager + BF16), which serves fmhaSm100f QkvBfloat16 ... DynamicTokenSparse cubins (B200 2026-10-04, + opharness finding sm100_trtllm_dsv4_facts_unservable_2026_10_04). Mirror that: the sm100 + collector lane is bf16-KV; the perf row carries kv_cache_dtype accordingly, spelled as the + SDK's KVCacheQuantMode member name ("bfloat16" / "fp8": engine_table_view rehydrates the + column through that enum, and "bf16" raised KeyError on every SM100 DSV4 table, Full CI + 2026-10-29). AIS_DSV4_KV_CACHE_DTYPE=fp8|bfloat16 (alias bf16) is the A/B hook. + """ + override = os.environ.get("AIS_DSV4_KV_CACHE_DTYPE") + if override: + return _KV_CACHE_DTYPE_LABELS.get(override, override) + return "fp8" if get_sm_version() in (90, 120, 121) else "bfloat16" + + def _build_dsv4_test_cases(mode: str, attn_kind: str) -> list[dict]: cases: list[dict] = [] tp_sizes = [1] if "--smoke" in sys.argv else _DSV4_MODULE_TP_SIZES @@ -179,7 +204,7 @@ def _build_dsv4_test_cases(mode: str, attn_kind: str) -> list[dict]: for model_path in model_paths: for tp_size in tp_sizes: for bs, sl, prefix in shapes: - params = [sl, bs, tp_size, "fp8", "bfloat16", "fp8_block", model_path, attn_kind] + params = [sl, bs, tp_size, _serving_kv_cache_dtype(), "bfloat16", "fp8_block", model_path, attn_kind] case_id = f"dsv4_{attn_kind}_{mode}_b{bs}_s{sl}_tp{tp_size}_{model_path.replace('/', '_')}" if mode == "context": params.append(prefix) @@ -222,10 +247,14 @@ def get_dsv4_hca_generation_test_cases() -> list[dict]: # untouched, and reusing one module across batches is exactly what serving # does. Owner-approved perf change 2026-08-09. _MODULE_CACHE: dict = {} +# Lower bound on the per-case KV token quota (see create_dsv4_kv_cache_and_metadata). +_KV_POOL_MIN_TOKENS = 2 * 1024 * 1024 -def _cached_dsv4_attention_module(model_path: str, attn_kind: str, tp_size: int, device: str): - key = (model_path, attn_kind, int(tp_size), device) +def _cached_dsv4_attention_module( + model_path: str, attn_kind: str, tp_size: int, device: str, kv_cache_dtype: str = "fp8" +): + key = (model_path, attn_kind, int(tp_size), device, kv_cache_dtype) hit = _MODULE_CACHE.get(key) if hit is not None: return hit @@ -235,7 +264,9 @@ def _cached_dsv4_attention_module(model_path: str, attn_kind: str, tp_size: int, _MODULE_CACHE.clear() gc.collect() torch.cuda.empty_cache() - entry = create_dsv4_attention_module(model_path=model_path, attn_kind=attn_kind, tp_size=tp_size, device=device) + entry = create_dsv4_attention_module( + model_path=model_path, attn_kind=attn_kind, tp_size=tp_size, device=device, kv_cache_dtype=kv_cache_dtype + ) _MODULE_CACHE[key] = entry return entry @@ -322,6 +353,7 @@ def create_dsv4_attention_module( attn_kind: str, tp_size: int, device: str = "cuda:0", + kv_cache_dtype: str = "fp8", ): """Build DeepseekV4Attention through TRT-LLM's own config path. @@ -379,7 +411,11 @@ def create_dsv4_attention_module( "The framework dispatch contract changed; re-audit against the runtime version." ) - _apply_gemm_type_quant(model_config, "fp8_block", use_fp8_kv_cache=True) + # The KV quant algo follows the lane: fp8 (Hopper fp8_ds_mla pool) sets kv_cache_quant_algo=FP8 and the + # module runs the E4m3 FMHA cubins; bf16 (the only servable SM100 pool, see _serving_kv_cache_dtype) + # leaves it unset so the module runs the QkvBfloat16 cubins serving runs — with FP8 pinned here the + # sm100 bf16 pool still produced E4m3 kernels (B200 capture 2026-10-04). + _apply_gemm_type_quant(model_config, "fp8_block", use_fp8_kv_cache=(kv_cache_dtype == "fp8")) # Provenance: print the RESOLVED kernel-selection knobs (same auto-build # code path default serving takes when the user sets no @@ -405,12 +441,43 @@ def create_dsv4_attention_module( ) ) + if get_sm_version() == 90: + # Serving publishes the requested pool dtype to the attention modules through the + # model's extra attrs (pyexecutor/model_loader.py:1934 extra_attrs['kv_cache_dtype'] + # = kv_cache_config.dtype; read by attention/mla.py:341). On Hopper that is + # 'fp8_ds_mla' (see the KvCacheConfig below) and it also switches the DSV4 module + # OFF the fused Q-FP8/KV-norm prologue (sparse/deepseek_v4/module.py:470-480 + # _is_fused_q_fp8_quant_enabled: "fp8_ds_mla does not use the fused Q FP8 path"). + # Without it the module takes the fused path, hands the footer-scale cache a + # strided raw latent and mla_rope_inplace rejects it ("data must be contiguous", + # H20 ctx capture 2026-10-01) — a path serving never runs on Hopper. + model_config.extra_attrs["kv_cache_dtype"] = "fp8_ds_mla" + aux_stream = torch.cuda.Stream(device=device) - attn_module = DeepseekV4Attention( - model_config=model_config, - layer_idx=0, - aux_stream=aux_stream, - ) + import inspect + + _ctor_params = inspect.signature(DeepseekV4Attention.__init__).parameters + if "aux_stream_dict" in _ctor_params: + # 1.3.0rc29: DeepseekV4Attention takes the model's per-purpose stream dict + # (modeling_deepseekv4.py:1371-1376, built by DeepseekV4ForCausalLM :1521-1570), + # as collect_mla_module already does for the MLA module. + from tensorrt_llm._torch.utils import AuxStreamType + + attn_module = DeepseekV4Attention( + model_config=model_config, + layer_idx=0, + aux_stream_dict={ + AuxStreamType.Attention: aux_stream, + AuxStreamType.MoeShared: aux_stream, + AuxStreamType.MoeChunkingOverlap: torch.cuda.Stream(device=device), + }, + ) + else: + attn_module = DeepseekV4Attention( + model_config=model_config, + layer_idx=0, + aux_stream=aux_stream, + ) # Serving applies QuantConfig.exclude_modules before weight creation # (apply_quant_config_exclude_modules, modeling_utils.py @runtime version); @@ -452,6 +519,7 @@ def create_dsv4_kv_cache_and_metadata( is_context: bool, prefix_len: int = 0, device: str = "cuda:0", + kv_cache_dtype_str: str = "fp8", ): """DSV4 cache manager + attention metadata, following the serving construction path with a DIRECT pinned citation on every hand-set field @@ -474,11 +542,20 @@ def create_dsv4_kv_cache_and_metadata( (is_gen: req.prompt_len = token_num - 1; py_prompt_len = prompt_len), consumed at model_engine.py:4336 """ - from tensorrt_llm._torch.attention_backend.interface import ( - AttentionRuntimeFeatures, - KVCacheParams, - ) - from tensorrt_llm._torch.attention_backend.utils import get_attention_backend + try: + from tensorrt_llm._torch.attention.backends.interface import ( + AttentionRuntimeFeatures, + KVCacheParams, + ) + except ModuleNotFoundError: # < rc29 layout + from tensorrt_llm._torch.attention_backend.interface import ( + AttentionRuntimeFeatures, + KVCacheParams, + ) + try: + from tensorrt_llm._torch.attention.backends.utils import get_attention_backend + except ModuleNotFoundError: # < rc29 layout + from tensorrt_llm._torch.attention_backend.utils import get_attention_backend from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal.batch_manager import CacheType @@ -537,20 +614,69 @@ def create_dsv4_kv_cache_and_metadata( # equivalent byte quota directly (half of currently-free device memory, # matching the DSV4 example's free_gpu_memory_fraction=0.5, # examples/models/core/deepseek_v4/README.md:148-151 @v1.3.0rc23). A - # main-KV-shaped max_tokens cap under-counts on V2 (the byte quota spans - # main KV + SWA + compressor/indexer caches) and made large-KV shapes - # fail dummy-request allocation ("Request ID not found in IndexMapper", - # B200 smoke round 1 2026-08-06). + # collector-computed main-KV-shaped byte cap under-counted on V2 (the + # byte quota spans main KV + SWA + compressor/indexer caches) and made + # large-KV shapes fail dummy-request allocation ("Request ID not found in + # IndexMapper", B200 smoke round 1 2026-08-06) — hence the token cap below + # is handed to the manager's own arithmetic, never computed here. free_bytes, _ = torch.cuda.mem_get_info(torch.device(device)) - kv_cache_config = KvCacheConfig( + # Pool size is a per-case cost, not a measurement input: the manager + # physically backs the whole quota at construction (runtime/ + # kv_cache_manager_v2/_storage_manager.py:248-249 slot counts from the + # GPU-tier quota, expand_pool_group :826 @1.3.0rc29) and this collector + # builds and tears one down per case. Cap the quota at what this batch + # can touch via kv_cache_config.max_tokens — the manager converts it with + # its own DSV4 arithmetic (DeepseekV4CacheManager._get_quota_from_max_tokens, + # sparse/deepseek_v4/cache_manager.py:926-981 @1.3.0rc29: non-sliding KV + + # SWA per token/request + indexer/compress pools + per-type padding) and + # takes min(max_gpu_total_bytes, quota_from_max_tokens) (pyexecutor/kv_cache/ + # kv_cache_manager_v2.py:1384-1398). The envelope is doubled and rounded + # to whole pages so dummy-request allocation never sits at the edge; a + # pool that is still too small fails allocation loudly (classified), + # never silently. The free-memory half stays as the upper bound, matching + # the DSV4 example's free_gpu_memory_fraction=0.5. + planned_tokens = batch_size * (max(max_seq, request_tokens) + tokens_per_block) + planned_tokens = -(-2 * planned_tokens // tokens_per_block) * tokens_per_block + # Floor: the pool must stay large enough that inductor emits the SAME + # kernels serving runs. footer_scale_kv.dequant_gather (fp8_ds_mla path, + # @maybe_compile(dynamic=True)) indexes the whole footer-scale pool; when + # that pool extent is < 2**31 bytes inductor generates int32-indexed + # triton kernels (xnumel i32), above it int64 — serving pools are tens of + # GiB, so serving always runs the int64 variant. H20 calibration + # (2026-10-03, DeepSeek-V4-Flash, b64 s128 p2048): quota from 1,048,576 + # tokens = 11.4 GiB -> i32 (53.9 ms), 1,572,864 tokens = 17.1 GiB -> i64 + # (45.8 ms); 2M tokens (22.7 GiB on Flash, i64 verified in the same sweep) + # sits 1.33x past the threshold, and every other DSV4 geometry carries more + # footer bytes per planned token than Flash, so the margin only grows. The + # footer row format is fixed (584 B/token), so the threshold is in tokens, + # not model geometry. 4M was tried first: on DeepSeek-V4-Pro it reproduced + # the campaign's free*0.5 pool size and left 20 GiB for the 262k-295k token + # context cells, whose forward allocates a 32 GiB intermediate -> OOM. + planned_tokens = max(planned_tokens, _KV_POOL_MIN_TOKENS) + kv_cache_kwargs = dict( tokens_per_block=tokens_per_block, max_gpu_total_bytes=int(free_bytes * 0.5), + max_tokens=int(planned_tokens), enable_block_reuse=False, ) + if get_sm_version() == 90: + # 1.3.0rc29 on Hopper: DeepseekV4CacheManager refuses any pool but the packed + # FP8 sparse-MLA cache (attention/backends/sparse/deepseek_v4/cache_manager.py: + # 343-349 "DeepSeek-V4 on Hopper requires kv_cache_config.dtype='fp8_ds_mla'"), + # the option llm_args.py:4256-4262 documents for SM90/SM120/SM121; serving + # deployments on H20 carry kv_cache_config.dtype=fp8_ds_mla (harness targets + # customization, matrix pass+custom). Earlier rc KvCacheConfig has no such + # literal and used the plain FP8 pool — keep that path for them. + try: + kv_cache_config = KvCacheConfig(dtype="fp8_ds_mla", **kv_cache_kwargs) + except (ValueError, TypeError): + kv_cache_config = KvCacheConfig(**kv_cache_kwargs) + else: + kv_cache_config = KvCacheConfig(**kv_cache_kwargs) kv_cache_manager_cls = get_kv_cache_manager_cls(model_config, kv_cache_config) # fp8 KV rows -> DataType.FP8, matching serving's kv_cache_dtype # resolution from kv_cache_quant_algo (set by _apply_gemm_type_quant). - kv_cache_dtype = DataType.FP8 + kv_cache_dtype = DataType.BF16 if kv_cache_dtype_str == "bfloat16" else DataType.FP8 # Serving construction site: pyexecutor/_util.py:1843-1867 @1.3.0rc23 # (is_mla branch): CacheType.SELFKONLY :1846, num_kv_heads=1 :1848, @@ -709,8 +835,8 @@ def run_dsv4_attn( if attn_kind not in ATTN_KIND_TO_COMPRESS_RATIO: raise ValueError(f"unsupported DSV4 attn_kind: {attn_kind}") - if kv_cache_dtype != "fp8": - raise ValueError(f"DSV4 module rows are fp8-KV only (got {kv_cache_dtype!r})") + if kv_cache_dtype not in ("fp8", "bfloat16"): + raise ValueError(f"DSV4 module rows are fp8-KV or bfloat16-KV only (got {kv_cache_dtype!r})") if compute_dtype != "bfloat16": raise ValueError(f"DSV4 module rows are bfloat16-compute only (got {compute_dtype!r})") if gemm_type != "fp8_block": @@ -720,7 +846,9 @@ def run_dsv4_attn( torch_device = torch.device(device) torch.cuda.set_device(torch_device) - attn_module, model_config, head_info = _cached_dsv4_attention_module(model_path, attn_kind, tp_size, device) + attn_module, model_config, head_info = _cached_dsv4_attention_module( + model_path, attn_kind, tp_size, device, kv_cache_dtype + ) # Ownership: the KV pool and the process-global extra-attrs slot are # released/restored on EVERY exit path (success, dry-run failure, @@ -738,6 +866,7 @@ def run_dsv4_attn( is_context=is_context, prefix_len=prefix_len, device=device, + kv_cache_dtype_str=kv_cache_dtype, ) hidden_size = model_config.pretrained_config.hidden_size diff --git a/python/aisimulate/collector/trtllm/collect_gdn.py b/python/aisimulate/collector/trtllm/collect_gdn.py index 5307edbe6a..240401ca45 100644 --- a/python/aisimulate/collector/trtllm/collect_gdn.py +++ b/python/aisimulate/collector/trtllm/collect_gdn.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -__compat__ = "trtllm>=1.3.0rc20" +__compat__ = "trtllm==1.3.0rc29" """ GDN (Gated DeltaNet) Collector for AIConfigurator — TensorRT-LLM backend. @@ -42,6 +42,8 @@ from tensorrt_llm._torch.modules.fla.fused_recurrent import fused_recurrent_gated_delta_rule from tensorrt_llm._torch.modules.mamba.causal_conv1d import causal_conv1d_fn, causal_conv1d_update + _CHUNK_KERNEL_SOURCE: str = "" # bound with the kernels by the runtime import (globals().update) + import torch try: @@ -167,6 +169,13 @@ def run_gdn_context_benchmark( conv_weight = torch.randn(conv_channels, d_conv, dtype=dtype, device=device) conv_bias = torch.randn(conv_channels, dtype=dtype, device=device) + # Every cell is timed under CUDA-graph replay; a failed capture raises + # (benchmark_with_power default allow_graph_fail=False) and the cell is + # counted below, so the case ends as a classified failure instead of an + # eager-timed row whose latency floor is the Python launch loop. + successful_points = 0 + failed_points = 0 + failures: list[str] = [] for batch_size in batch_size_list: for seq_len in seq_len_list: if aic_debug: @@ -199,10 +208,22 @@ def run_gdn_context_benchmark( q = torch.randn(batch_size, seq_len, num_k_heads, head_k_dim, dtype=dtype, device=device) k = torch.randn(batch_size, seq_len, num_k_heads, head_k_dim, dtype=dtype, device=device) v = torch.randn(batch_size, seq_len, num_v_heads, head_v_dim, dtype=dtype, device=device) + # g is fp32 in serving (fused_gdn_gating output); FlashInfer's + # chunk kernel asserts it ("g must be fp32"), Triton accepts both g = torch.nn.functional.logsigmoid( - torch.randn(batch_size, seq_len, num_v_heads, dtype=dtype, device=device) + torch.randn(batch_size, seq_len, num_v_heads, dtype=torch.float32, device=device) ) beta = torch.sigmoid(torch.randn(batch_size, seq_len, num_v_heads, dtype=dtype, device=device)) + # serving feeds the chunk kernel PACKED tokens [1, B*S, H, D] plus + # cu_seqlens (gdn_mixer.forward_extend); FlashInfer's chunk kernel + # accepts only that layout ("q must be [1, T, H_q, D_k]") + q, k, v, g, beta = (t.reshape(1, batch_size * seq_len, *t.shape[2:]) for t in (q, k, v, g, beta)) + cu_seqlens = torch.arange(0, batch_size * seq_len + 1, seq_len, dtype=torch.int64, device=device) + # serving passes the per-request recurrent state and asks for the + # final state back (gdn_mixer.py:847-855); FlashInfer requires it + gdn_init_state = torch.zeros( + batch_size, num_v_heads, head_k_dim, head_v_dim, dtype=torch.bfloat16, device=device + ) # --- Benchmark causal_conv1d_fn --- torch.cuda.synchronize() @@ -218,7 +239,6 @@ def run_conv1d(_ki=k_input, _cs=conv_state): num_warmups=num_warmups, num_runs=num_runs, repeat_n=1, - allow_graph_fail=True, ) as results: log_perf( item_list=[{**common_log_data, "latency": results["latency_ms"]}], @@ -233,11 +253,15 @@ def run_conv1d(_ki=k_input, _cs=conv_state): # --- Benchmark chunk_gated_delta_rule --- torch.cuda.synchronize() - chunk_gated_delta_rule(q, k, v, g, beta) + chunk_gated_delta_rule( + q, k, v, g, beta, initial_state=gdn_init_state, output_final_state=True, cu_seqlens=cu_seqlens + ) torch.cuda.synchronize() - def run_gdn_scan(_q=q, _k=k, _v=v, _g=g, _beta=beta): - chunk_gated_delta_rule(_q, _k, _v, _g, _beta) + def run_gdn_scan(_q=q, _k=k, _v=v, _g=g, _beta=beta, _cu=cu_seqlens, _st=gdn_init_state): + chunk_gated_delta_rule( + _q, _k, _v, _g, _beta, initial_state=_st, output_final_state=True, cu_seqlens=_cu + ) with benchmark_with_power( device=device, @@ -245,7 +269,6 @@ def run_gdn_scan(_q=q, _k=k, _v=v, _g=g, _beta=beta): num_warmups=num_warmups, num_runs=num_runs, repeat_n=1, - allow_graph_fail=True, ) as results: log_perf( item_list=[{**common_log_data, "latency": results["latency_ms"]}], @@ -253,7 +276,7 @@ def run_gdn_scan(_q=q, _k=k, _v=v, _g=g, _beta=beta): version=trtllm_version, device_name=torch.cuda.get_device_name(device), op_name="gdn", - kernel_source="chunk_gated_delta_rule", + kernel_source=_CHUNK_KERNEL_SOURCE, perf_filename=perf_filename, power_stats=results["power_stats"], ) @@ -262,16 +285,22 @@ def run_gdn_scan(_q=q, _k=k, _v=v, _g=g, _beta=beta): input_pool = _make_input_pool( { "k_input": (batch_size, conv_channels, seq_len), - "q": (batch_size, seq_len, num_k_heads, head_k_dim), - "k": (batch_size, seq_len, num_k_heads, head_k_dim), - "v": (batch_size, seq_len, num_v_heads, head_v_dim), - "g": (batch_size, seq_len, num_v_heads), - "beta": (batch_size, seq_len, num_v_heads), + # packed [1, B*S, ...] + cu_seqlens, as serving feeds the chunk kernel + "q": (1, batch_size * seq_len, num_k_heads, head_k_dim), + "k": (1, batch_size * seq_len, num_k_heads, head_k_dim), + "v": (1, batch_size * seq_len, num_v_heads, head_v_dim), + "g": (1, batch_size * seq_len, num_v_heads), + "beta": (1, batch_size * seq_len, num_v_heads), }, total_iters, dtype, device, ) + input_pool["g"] = [t.float() for t in input_pool["g"]] # serving g is fp32 + cu_seqlens = torch.arange(0, batch_size * seq_len + 1, seq_len, dtype=torch.int64, device=device) + gdn_init_state = torch.zeros( + batch_size, num_v_heads, head_k_dim, head_v_dim, dtype=torch.bfloat16, device=device + ) for i in range(total_iters): input_pool["g"][i] = torch.nn.functional.logsigmoid(input_pool["g"][i]) input_pool["beta"][i] = torch.sigmoid(input_pool["beta"][i]) @@ -298,7 +327,6 @@ def run_conv1d(_pool=input_pool, _cs=conv_state, _idx=conv1d_iter_idx): num_warmups=num_warmups, num_runs=num_runs, repeat_n=1, - allow_graph_fail=True, ) as results: log_perf( item_list=[{**common_log_data, "latency": results["latency_ms"]}], @@ -319,6 +347,9 @@ def run_conv1d(_pool=input_pool, _cs=conv_state, _idx=conv1d_iter_idx): input_pool["v"][0], input_pool["g"][0], input_pool["beta"][0], + initial_state=gdn_init_state, + output_final_state=True, + cu_seqlens=cu_seqlens, ) torch.cuda.synchronize() @@ -333,6 +364,9 @@ def run_gdn_scan(_pool=input_pool, _idx=gdn_iter_idx): _pool["v"][idx], _pool["g"][idx], _pool["beta"][idx], + initial_state=gdn_init_state, + output_final_state=True, + cu_seqlens=cu_seqlens, ) with benchmark_with_power( @@ -341,7 +375,6 @@ def run_gdn_scan(_pool=input_pool, _idx=gdn_iter_idx): num_warmups=num_warmups, num_runs=num_runs, repeat_n=1, - allow_graph_fail=True, ) as results: log_perf( item_list=[{**common_log_data, "latency": results["latency_ms"]}], @@ -349,7 +382,7 @@ def run_gdn_scan(_pool=input_pool, _idx=gdn_iter_idx): version=trtllm_version, device_name=torch.cuda.get_device_name(device), op_name="gdn", - kernel_source="chunk_gated_delta_rule", + kernel_source=_CHUNK_KERNEL_SOURCE, perf_filename=perf_filename, power_stats=results["power_stats"], ) @@ -361,11 +394,108 @@ def run_gdn_scan(_pool=input_pool, _idx=gdn_iter_idx): del input_pool, conv_state gc.collect() torch.cuda.empty_cache() + successful_points += 1 except Exception as e: + failed_points += 1 + failures.append(f"batch_size={batch_size}, seq_len={seq_len}: {type(e).__name__}: {e}") print(f" Error at batch_size={batch_size}, seq_len={seq_len}: {e}") continue + summary = f"ok={successful_points} error={failed_points} skip=0" + print(f"GDN context summary: {summary}") + if failed_points or successful_points == 0: + raise RuntimeError( + f"TRTLLM GDN context collection failed strict completeness: {summary}; failed cells: {'; '.join(failures)}" + ) + + +def _bench_serving_decode_entry( + *, + batch_size: int, + num_k_heads: int, + head_k_dim: int, + num_v_heads: int, + head_v_dim: int, + dtype, + device, + common_log_data: dict, + trtllm_version: str, + perf_filename: str, + num_warmups: int, + num_runs: int, +) -> None: + """Measure the entry serving actually calls for pure decode. + + gdn_mixer.forward_decode (_torch/modules/mamba/gdn_mixer.py:670 @1.3.0rc23) + calls fused_sigmoid_gating_delta_rule_update with the raw gating inputs + (A_log, dt_bias, a, b), the SSM state pool + slot indices and + cu_seqlens; that function dispatches to FlashInfer's bf16-state GDN decode + kernel when _can_use_flashinfer_gdn_decode holds (SM90/SM100, K==V==128, + bf16 state, one token per sequence) and to the FLA Triton kernel otherwise + (fla/fused_sigmoid_gating_recurrent.py:205-238,452-458). kernel_source + records which branch ran. Builds without this entry (pre-rc23) skip it — + their serving truth is the fused_recurrent row above. + """ + fn = globals().get("fused_sigmoid_gating_delta_rule_update") + gate = globals().get("_can_use_flashinfer_gdn_decode") + if fn is None: + return + B, H, K, V = batch_size, num_v_heads, head_k_dim, head_v_dim # noqa: N806 + q = torch.randn(1, B, num_k_heads, K, dtype=dtype, device=device) + k = torch.randn(1, B, num_k_heads, K, dtype=dtype, device=device) + v = torch.randn(1, B, H, V, dtype=dtype, device=device) + a = torch.randn(B, H, dtype=dtype, device=device) + b = torch.randn(B, H, dtype=dtype, device=device) + A_log = torch.zeros(H, dtype=torch.float32, device=device) # noqa: N806 + dt_bias = 0.1 * torch.randn(H, dtype=torch.float32, device=device) + ssm_states = torch.randn(B, H, K, V, dtype=torch.bfloat16, device=device) + slot_indices = torch.arange(B, dtype=torch.int32, device=device) + cu_seqlens = torch.arange(B + 1, dtype=torch.int64, device=device) + uses_flashinfer = bool(gate(ssm_states, K, V, B, B)) if gate is not None else False + + def run_serving_decode(): + fn( + A_log=A_log, + dt_bias=dt_bias, + q=q, + k=k, + v=v, + a=a, + b=b, + initial_state_source=ssm_states, + initial_state_indices=slot_indices, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=True, + softplus_beta=1.0, + softplus_threshold=20.0, + ) + + torch.cuda.synchronize() + run_serving_decode() + torch.cuda.synchronize() + with benchmark_with_power( + device=device, + kernel_func=run_serving_decode, + num_warmups=num_warmups, + num_runs=num_runs, + repeat_n=1, + ) as results: + log_perf( + item_list=[{**common_log_data, "latency": results["latency_ms"]}], + framework="TRTLLM", + version=trtllm_version, + device_name=torch.cuda.get_device_name(device), + op_name="gdn", + kernel_source=( + "fused_sigmoid_gating_delta_rule_update_flashinfer" + if uses_flashinfer + else "fused_sigmoid_gating_delta_rule_update_triton" + ), + perf_filename=perf_filename, + power_stats=results["power_stats"], + ) + def run_gdn_generation_benchmark( d_model: int, @@ -404,6 +534,11 @@ def run_gdn_generation_benchmark( conv_weight = torch.randn(conv_channels, d_conv, dtype=dtype, device=device) conv_bias = torch.randn(conv_channels, dtype=dtype, device=device) + # CUDA-graph replay only (see run_gdn_context_benchmark): a failed capture + # raises and the batch is counted as a failed cell below. + successful_points = 0 + failed_points = 0 + failures: list[str] = [] for batch_size in batch_size_list: if aic_debug: print(f" Benchmarking batch_size={batch_size}") @@ -431,12 +566,31 @@ def run_gdn_generation_benchmark( "model_name": model_name, } + # serving's pure-decode entry (dispatch-labeled); the rows below keep + # the conv update and the FLA recurrent fallback lane + _bench_serving_decode_entry( + batch_size=batch_size, + num_k_heads=num_k_heads, + head_k_dim=head_k_dim, + num_v_heads=num_v_heads, + head_v_dim=head_v_dim, + dtype=dtype, + device=device, + common_log_data=common_log_data, + trtllm_version=trtllm_version, + perf_filename=perf_filename, + num_warmups=num_warmups, + num_runs=num_runs, + ) + if aic_cached_inputs: k_input = torch.randn(batch_size, conv_channels, dtype=dtype, device=device) q = torch.randn(batch_size, 1, num_k_heads, head_k_dim, dtype=dtype, device=device) k = torch.randn(batch_size, 1, num_k_heads, head_k_dim, dtype=dtype, device=device) v = torch.randn(batch_size, 1, num_v_heads, head_v_dim, dtype=dtype, device=device) - g = torch.nn.functional.logsigmoid(torch.randn(batch_size, 1, num_v_heads, dtype=dtype, device=device)) + g = torch.nn.functional.logsigmoid( + torch.randn(batch_size, 1, num_v_heads, dtype=torch.float32, device=device) + ) # serving g is fp32 beta = torch.sigmoid(torch.randn(batch_size, 1, num_v_heads, dtype=dtype, device=device)) # --- Benchmark causal_conv1d_update --- @@ -453,7 +607,6 @@ def run_conv1d_update(_k=k_input, _cs=conv_state): num_warmups=num_warmups, num_runs=num_runs, repeat_n=1, - allow_graph_fail=True, ) as results: log_perf( item_list=[{**common_log_data, "latency": results["latency_ms"]}], @@ -466,48 +619,52 @@ def run_conv1d_update(_k=k_input, _cs=conv_state): power_stats=results["power_stats"], ) - # --- Benchmark fused_recurrent_gated_delta_rule --- - torch.cuda.synchronize() - fused_recurrent_gated_delta_rule( - q, - k, - v, - g, - beta, - initial_state=gdn_state, - output_final_state=True, - ) - torch.cuda.synchronize() - - def run_gdn_update(_q=q, _k=k, _v=v, _g=g, _beta=beta, _state=gdn_state): + if fused_recurrent_gated_delta_rule is None: + print( + "[gdn] fused_recurrent_gated_delta_rule row omitted: kernel absent in this trtllm (>=1.3.0rc29)" + ) + else: + # --- Benchmark fused_recurrent_gated_delta_rule --- + torch.cuda.synchronize() fused_recurrent_gated_delta_rule( - _q, - _k, - _v, - _g, - _beta, - initial_state=_state, + q, + k, + v, + g, + beta, + initial_state=gdn_state, output_final_state=True, ) + torch.cuda.synchronize() - with benchmark_with_power( - device=device, - kernel_func=run_gdn_update, - num_warmups=num_warmups, - num_runs=num_runs, - repeat_n=1, - allow_graph_fail=True, - ) as results: - log_perf( - item_list=[{**common_log_data, "latency": results["latency_ms"]}], - framework="TRTLLM", - version=trtllm_version, - device_name=torch.cuda.get_device_name(device), - op_name="gdn", - kernel_source="fused_recurrent_gated_delta_rule", - perf_filename=perf_filename, - power_stats=results["power_stats"], - ) + def run_gdn_update(_q=q, _k=k, _v=v, _g=g, _beta=beta, _state=gdn_state): + fused_recurrent_gated_delta_rule( + _q, + _k, + _v, + _g, + _beta, + initial_state=_state, + output_final_state=True, + ) + + with benchmark_with_power( + device=device, + kernel_func=run_gdn_update, + num_warmups=num_warmups, + num_runs=num_runs, + repeat_n=1, + ) as results: + log_perf( + item_list=[{**common_log_data, "latency": results["latency_ms"]}], + framework="TRTLLM", + version=trtllm_version, + device_name=torch.cuda.get_device_name(device), + op_name="gdn", + kernel_source="fused_recurrent_gated_delta_rule", + perf_filename=perf_filename, + power_stats=results["power_stats"], + ) else: input_pool = _make_input_pool( @@ -523,6 +680,7 @@ def run_gdn_update(_q=q, _k=k, _v=v, _g=g, _beta=beta, _state=gdn_state): dtype, device, ) + input_pool["g"] = [t.float() for t in input_pool["g"]] # serving g is fp32 for i in range(total_iters): input_pool["g"][i] = torch.nn.functional.logsigmoid(input_pool["g"][i]) input_pool["beta"][i] = torch.sigmoid(input_pool["beta"][i]) @@ -545,7 +703,6 @@ def run_conv1d_update(_pool=input_pool, _cs=conv_state, _idx=conv1d_iter_idx): num_warmups=num_warmups, num_runs=num_runs, repeat_n=1, - allow_graph_fail=True, ) as results: log_perf( item_list=[{**common_log_data, "latency": results["latency_ms"]}], @@ -558,52 +715,56 @@ def run_conv1d_update(_pool=input_pool, _cs=conv_state, _idx=conv1d_iter_idx): power_stats=results["power_stats"], ) - # --- Benchmark fused_recurrent_gated_delta_rule --- - torch.cuda.synchronize() - fused_recurrent_gated_delta_rule( - input_pool["q"][0], - input_pool["k"][0], - input_pool["v"][0], - input_pool["g"][0], - input_pool["beta"][0], - initial_state=gdn_state, - output_final_state=True, - ) - torch.cuda.synchronize() - - gdn_iter_idx = [0] - - def run_gdn_update(_pool=input_pool, _state=gdn_state, _idx=gdn_iter_idx): - idx = _idx[0] % total_iters - _idx[0] += 1 + if fused_recurrent_gated_delta_rule is None: + print( + "[gdn] fused_recurrent_gated_delta_rule row omitted: kernel absent in this trtllm (>=1.3.0rc29)" + ) + else: + # --- Benchmark fused_recurrent_gated_delta_rule --- + torch.cuda.synchronize() fused_recurrent_gated_delta_rule( - _pool["q"][idx], - _pool["k"][idx], - _pool["v"][idx], - _pool["g"][idx], - _pool["beta"][idx], - initial_state=_state, + input_pool["q"][0], + input_pool["k"][0], + input_pool["v"][0], + input_pool["g"][0], + input_pool["beta"][0], + initial_state=gdn_state, output_final_state=True, ) + torch.cuda.synchronize() - with benchmark_with_power( - device=device, - kernel_func=run_gdn_update, - num_warmups=num_warmups, - num_runs=num_runs, - repeat_n=1, - allow_graph_fail=True, - ) as results: - log_perf( - item_list=[{**common_log_data, "latency": results["latency_ms"]}], - framework="TRTLLM", - version=trtllm_version, - device_name=torch.cuda.get_device_name(device), - op_name="gdn", - kernel_source="fused_recurrent_gated_delta_rule", - perf_filename=perf_filename, - power_stats=results["power_stats"], - ) + gdn_iter_idx = [0] + + def run_gdn_update(_pool=input_pool, _state=gdn_state, _idx=gdn_iter_idx): + idx = _idx[0] % total_iters + _idx[0] += 1 + fused_recurrent_gated_delta_rule( + _pool["q"][idx], + _pool["k"][idx], + _pool["v"][idx], + _pool["g"][idx], + _pool["beta"][idx], + initial_state=_state, + output_final_state=True, + ) + + with benchmark_with_power( + device=device, + kernel_func=run_gdn_update, + num_warmups=num_warmups, + num_runs=num_runs, + repeat_n=1, + ) as results: + log_perf( + item_list=[{**common_log_data, "latency": results["latency_ms"]}], + framework="TRTLLM", + version=trtllm_version, + device_name=torch.cuda.get_device_name(device), + op_name="gdn", + kernel_source="fused_recurrent_gated_delta_rule", + perf_filename=perf_filename, + power_stats=results["power_stats"], + ) # Cleanup if aic_cached_inputs: @@ -612,11 +773,22 @@ def run_gdn_update(_pool=input_pool, _state=gdn_state, _idx=gdn_iter_idx): del input_pool, conv_state, gdn_state gc.collect() torch.cuda.empty_cache() + successful_points += 1 except Exception as e: + failed_points += 1 + failures.append(f"batch_size={batch_size}: {type(e).__name__}: {e}") print(f" Error at batch_size={batch_size}: {e}") continue + summary = f"ok={successful_points} error={failed_points} skip=0" + print(f"GDN generation summary: {summary}") + if failed_points or successful_points == 0: + raise RuntimeError( + f"TRTLLM GDN generation collection failed strict completeness: {summary}; " + f"failed cells: {'; '.join(failures)}" + ) + def run_gdn_torch( phase: str, @@ -647,17 +819,60 @@ def run_gdn_torch( contextlib.redirect_stderr(_devnull_file), ): import tensorrt_llm - from tensorrt_llm._torch.modules.fla.chunk import chunk_gated_delta_rule - from tensorrt_llm._torch.modules.fla.fused_recurrent import fused_recurrent_gated_delta_rule from tensorrt_llm._torch.modules.mamba.causal_conv1d import causal_conv1d_fn, causal_conv1d_update + # Extra FLA row (the vendored single-step recurrence). trtllm >=1.3.0rc29 + # removed fused_recurrent_gated_delta_rule from fla/fused_recurrent.py + # (only the *_update variants remain; serving never called the removed + # one — gdn_mixer.forward_decode uses fused_sigmoid_gating_delta_rule_update + # at rc23:670 and rc29:706 alike). When the kernel does not exist the row + # cannot exist: the serving-dispatch row below is still produced, and the + # omission is printed once per case, never silent. + try: + from tensorrt_llm._torch.modules.fla.fused_recurrent import fused_recurrent_gated_delta_rule + except ImportError: + fused_recurrent_gated_delta_rule = None + + # Serving dispatch, never the vendored kernel directly: the GDN mixer + # resolves chunk_gated_delta_rule to FlashInfer's cute-DSL chunk kernel on + # SM90/SM100 (TLLM_USE_FLASHINFER_GDN_PREFILL=1 default) and to the FLA + # Triton kernel elsewhere (_torch/modules/mamba/gdn_mixer.py:57-70 + # @1.3.0rc23; arch gate _utils.is_flashinfer_gdn_supported_arch). Older + # builds have no resolver and only the Triton kernel — import that. + chunk_label = "chunk_gated_delta_rule" + try: + from tensorrt_llm._torch.modules.mamba.gdn_mixer import ( + _resolve_chunk_gated_delta_rule, + chunk_gated_delta_rule, + ) + + if "flashinfer" in getattr(_resolve_chunk_gated_delta_rule(), "__module__", ""): + chunk_label = "chunk_gated_delta_rule_flashinfer" + except ImportError: + from tensorrt_llm._torch.modules.fla.chunk import chunk_gated_delta_rule + # Pure-decode serving entry (gdn_mixer.py:670): fused sigmoid gating + + # recurrent update, itself dispatching to FlashInfer's bf16-state GDN + # decode kernel when _can_use_flashinfer_gdn_decode holds + # (fla/fused_sigmoid_gating_recurrent.py:205-238,452-458 @1.3.0rc23). + try: + from tensorrt_llm._torch.modules.fla.fused_sigmoid_gating_recurrent import ( + _can_use_flashinfer_gdn_decode, + fused_sigmoid_gating_delta_rule_update, + ) + except ImportError: + _can_use_flashinfer_gdn_decode = None + fused_sigmoid_gating_delta_rule_update = None + globals().update( { "tensorrt_llm": tensorrt_llm, "causal_conv1d_fn": causal_conv1d_fn, "causal_conv1d_update": causal_conv1d_update, "chunk_gated_delta_rule": chunk_gated_delta_rule, + "_CHUNK_KERNEL_SOURCE": chunk_label, "fused_recurrent_gated_delta_rule": fused_recurrent_gated_delta_rule, + "fused_sigmoid_gating_delta_rule_update": fused_sigmoid_gating_delta_rule_update, + "_can_use_flashinfer_gdn_decode": _can_use_flashinfer_gdn_decode, } ) diff --git a/python/aisimulate/collector/trtllm/collect_gemm.py b/python/aisimulate/collector/trtllm/collect_gemm.py index 7053d953a7..c65b3fabae 100644 --- a/python/aisimulate/collector/trtllm/collect_gemm.py +++ b/python/aisimulate/collector/trtllm/collect_gemm.py @@ -9,7 +9,7 @@ quantized-weight setup, and perf logging. """ -__compat__ = "trtllm>=1.3.0rc20" +__compat__ = "trtllm==1.3.0rc29" import ctypes import math @@ -19,13 +19,12 @@ import tensorrt_llm import torch import torch.nn.functional as F -from case_generator import get_gemm_case_specs +from collector.case_generator import get_gemm_case_specs +from collector.helper import benchmark_with_power, get_sm_version, log_perf from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._utils import is_sm_100f from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig -from helper import benchmark_with_power, get_sm_version, log_perf - def pad_up(x: int, y: int) -> int: return ((x + y - 1) // y) * y @@ -50,7 +49,47 @@ def _compute_fp8_block_weight_scales(weight: torch.Tensor, group_size: int) -> t # Per-worker-process weight cache for large shapes (outside_loop_count == 1). # Keyed by (gemm_type, n, k). Only the most recently used entry is kept so # that GPU memory is not accumulated across (n, k) groups. -_weight_cache: dict = {} +# One worker builds the Linear modules (and their weights) once per +# (gemm_type, n, k) and reuses them for every m that follows on this worker: +# the modules do not depend on m, and the per-case rebuild was the shard's +# cost — nvfp4 unswizzles its block scales on the CPU and the sub-L2 shapes +# build up to 5 copies per case (b300 base run, job 471159002: gemm/nvfp4 +# 1.3 GPU-s/row and gemm/fp8_block 1.7 GPU-s/row against 0.04 for bf16). +# get_gemm_test_cases orders the plan by (n, k) then gemm_type then m, so a +# worker sees runs of equal keys. The cache holds one key at a time, so peak +# memory is unchanged. +_gemm_op_cache: dict = {} + + +def _gemm_ops(gemm_type: str, n: int, k: int, *, quant_config, group_size, outside_loop_count: int, device, dtype, x): + """The ``outside_loop_count`` Linear copies for (gemm_type, n, k), built on first use.""" + key = (gemm_type, n, k, outside_loop_count) + op_list = _gemm_op_cache.get(key) + if op_list is not None: + return op_list + _gemm_op_cache.clear() # release the previous key's modules before building + torch.cuda.empty_cache() + op_list = [] + for _i in range(outside_loop_count): + # nvfp4's per-tensor activation scale is a calibration constant in + # serving; the first x of this key provides it, and the kernel's + # timing does not depend on its value. + weights = _build_weights(gemm_type, n, k, device, dtype, group_size, x) + gemm = Linear( + k, + n, + bias=False, + dtype=dtype, + quant_config=quant_config, + force_dynamic_quantization=False, + ) + gemm.load_weights([weights]) + if gemm_type == "fp8_block" and callable(getattr(gemm, "post_load_weights", None)): + gemm.post_load_weights() + gemm.to(torch.device(device)) + op_list.append(gemm) + _gemm_op_cache[key] = op_list + return op_list def _get_l2_cache_bytes(device_id: int = 0) -> int: @@ -167,7 +206,12 @@ def run_gemm(gemm_type, m, n, k, *, perf_filename, device="cuda:0"): group_size = 128 qc = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=group_size) elif gemm_type == "nvfp4": - group_size = 128 + # NVFP4 scale blocks are 16 elements (the weights above are quantized with + # fp4_quantize(..., 16, ...)); 1.3.0rc29 NVFP4LinearMethod.resolve_scaling_vector_size + # rejects any other declared group_size ("supports NVFP4 scale blocks of (16,) elements, + # but the checkpoint declares group_size=128" — first seen on B200 2026-10-04; on Hopper + # the nvfp4 lane is a platform floor and never constructed the module). + group_size = 16 qc = QuantConfig(quant_algo=QuantAlgo.NVFP4, group_size=group_size) else: qc = None @@ -177,63 +221,20 @@ def run_gemm(gemm_type, m, n, k, *, perf_filename, device="cuda:0"): _bytes_per_elem = {"bfloat16": 2, "fp8": 1, "fp8_block": 1, "nvfp4": 0.5} _weight_bytes = int(n * k * _bytes_per_elem[gemm_type]) outside_loop_count = max(1, min(5, math.ceil(_l2_cache_bytes / _weight_bytes))) - op_list = [] + op_list = _gemm_ops( + gemm_type, + n, + k, + quant_config=qc, + group_size=group_size, + outside_loop_count=outside_loop_count, + device=device, + dtype=dtype, + x=x, + ) + for op in op_list: + op.forward(x) # dry run at this m (lazy init / JIT) before graph capture - if outside_loop_count == 1: - # Large weight (exceeds GPU L2): reuse cached tensor across consecutive - # tasks with the same (gemm_type, n, k) to avoid re-running the slow - # torch.randn call for every batch-size variant. - cache_key = (gemm_type, n, k) - if cache_key not in _weight_cache: - _weight_cache.clear() # release the previous (n,k) weight - _weight_cache[cache_key] = _build_weights(gemm_type, n, k, device, dtype, group_size, x) - - weights = _weight_cache[cache_key] - - # For nvfp4 the input_scale is x-dependent; recompute it cheaply. - if gemm_type == "nvfp4": - x_sf_global = (448 * 6) / x.abs().max().float() - weights = {**weights, "input_scale": 1.0 / x_sf_global.cpu()} - - # The fp8 weights carry no input_scale, so TRT-LLM's FP8QDQ method runs - # its dynamic-quantization path and the collected fp8 GEMM includes the - # compute-scale overhead. AIC derives fp8_static at query time by - # subtracting that overhead, so no separate fp8_static GEMM is collected. - gemm = Linear( - k, - n, - bias=False, - dtype=dtype, - quant_config=qc, - force_dynamic_quantization=False, - ) - gemm.load_weights([weights]) - if gemm_type == "fp8_block" and callable(getattr(gemm, "post_load_weights", None)): - gemm.post_load_weights() - gemm.to(torch.device(device)) - gemm.forward(x) # dry run to init - op_list.append(gemm) - else: - # Small weight: create distinct copies so that repeated calls span more - # than the GPU L2 cache, giving a realistic memory-pressure measurement. - for _i in range(outside_loop_count): - weights = _build_weights(gemm_type, n, k, device, dtype, group_size, x) - gemm = Linear( - k, - n, - bias=False, - dtype=dtype, - quant_config=qc, - force_dynamic_quantization=False, - ) - gemm.load_weights([weights]) - if gemm_type == "fp8_block" and callable(getattr(gemm, "post_load_weights", None)): - gemm.post_load_weights() - gemm.to(torch.device(device)) - gemm.forward(x) # dry run to init - op_list.append(gemm) - - # Use benchmark_with_power context manager def kernel_func(): for op in op_list: op.forward(x) diff --git a/python/aisimulate/collector/trtllm/collect_mamba2.py b/python/aisimulate/collector/trtllm/collect_mamba2.py index c0bbe59073..4ef76599bf 100644 --- a/python/aisimulate/collector/trtllm/collect_mamba2.py +++ b/python/aisimulate/collector/trtllm/collect_mamba2.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -__compat__ = "trtllm>=1.3.0rc20" +__compat__ = "trtllm==1.3.0rc29" """ Mamba2 SSM Collector for AIConfigurator. @@ -125,6 +125,117 @@ def get_mamba2_test_cases(): return test_cases +def _token_major_prefill_conv() -> bool: + """Serving's prefill conv layout at this runtime. + + 1.3.0rc29 Mamba2Mixer defaults TRTLLM_MAMBA_TOKEN_MAJOR_CONV=1 (mamba2_mixer.py:319) + and feeds causal_conv1d_fn a channel-last VIEW of the token-major projection with a + separate token-major ``out`` (mamba2_mixer.py:541-569; causal_conv1d.py:26-60 picks + causal_conv1d_channellast_fwd_kernel off x.stride(0) == 1). Older runtimes have no + ``out`` parameter and run the channel-major kernel on (batch, dim, seq) inputs. The + H20 mamba2_ctx gate (2026-10-02) diverged on exactly this kernel pair. + """ + import inspect + + return ( + "out" in inspect.signature(causal_conv1d_fn).parameters + and os.environ.get("TRTLLM_MAMBA_TOKEN_MAJOR_CONV", "1") == "1" + ) + + +def _prefill_conv_token_major_inputs(xbc_batched, batch_size, seq_len, device): + """(batch, dim, seq) random input -> the serving varlen token-major buffer + [tokens, dim] (sequences concatenated), its ``out`` twin and the varlen + metadata serving passes (mamba2_mixer.py:536-569: cu_seqlens int32, + cache_indices, has_initial_state=False for a cache-cold prefill).""" + conv_dim = xbc_batched.shape[1] + x_tm = xbc_batched.permute(0, 2, 1).reshape(batch_size * seq_len, conv_dim).contiguous() + out_tm = torch.empty_like(x_tm) + cu_seqlens = torch.arange(0, batch_size * seq_len + 1, seq_len, dtype=torch.int32, device=device) + cache_indices = torch.arange(batch_size, dtype=torch.int32, device=device) + has_initial_state = torch.zeros(batch_size, dtype=torch.bool, device=device) + return x_tm, out_tm, cu_seqlens, cache_indices, has_initial_state + + +def _run_prefill_conv_token_major( + x_tm, out_tm, conv_weight, conv_bias, conv_state, cu_seqlens, cache_indices, has_initial_state +): + causal_conv1d_fn( + x_tm.t(), + conv_weight, + conv_bias, + query_start_loc=cu_seqlens, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + conv_states=conv_state, + activation="silu", + out=out_tm.t(), + ) + + +# Both SSM kernels take the output buffer from the caller: serving hands them a +# view of one preallocated per-forward buffer (mamba2_mixer.py:453-457 alloc, +# :634 prefill `out=`, :906 decode `out=` @1.3.0rc29; rc20 :335/:403/:629) and +# neither kernel allocates one when it is missing — ssd_chunk_scan.py:667 +# `assert out.shape == x.shape` and selective_state_update.py `out.dim()` +# (rc20 :85, rc29 same) both fail on None. The collector omitted `out` until +# 2026-10-12, so every scan/state-update cell raised and the per-cell except +# dropped the lane: no shipped mamba2 table has those rows. + +# Decode SSM state update: serving picks FlashInfer's kernel or TRT-LLM's own +# Triton kernel per layer geometry — Mamba2Mixer.__init__ (mamba2_mixer.py: +# 284-293 gate, :322-328 selection @1.3.0rc29; rc20 mixer :14-23 had the ratio +# list [1, 8, 16]). The collector measures the full (unsharded) layer, i.e. the +# TP=1 geometry the gate sees as tp_nheads / tp_ngroups. +# FIXME(re-verify): re-check the three lists on every TRT-LLM version bump. +_FLASHINFER_SSU_HEAD_DIMS = (64, 128) +_FLASHINFER_SSU_HEAD_GROUP_RATIOS = (1, 2, 4, 8, 16, 32, 64) +_FLASHINFER_SSU_D_STATES = (64, 128, 256) +SSU_KERNEL_SOURCE_FLASHINFER = "selective_state_update_flashinfer" +SSU_KERNEL_SOURCE_TRITON = "selective_state_update_triton" + + +def _serving_selects_flashinfer_ssu(nheads: int, n_groups: int, head_dim: int, d_state: int) -> bool: + head_group_ratio = nheads // n_groups if n_groups > 0 else 0 + return ( + head_dim in _FLASHINFER_SSU_HEAD_DIMS + and head_group_ratio in _FLASHINFER_SSU_HEAD_GROUP_RATIOS + and d_state in _FLASHINFER_SSU_D_STATES + ) + + +def _decode_state_update_kernel(nheads: int, n_groups: int, head_dim: int, d_state: int): + """The (callable, kernel_source) serving's Mamba2Mixer runs for this layer's decode. + + kernel_source records the function serving calls (`selective_state_update_fi` + from flashinfer.mamba, or TRT-LLM's own Triton `selective_state_update`), the + way the GDN tables record `chunk_gated_delta_rule_{flashinfer,triton}`. No + fallback: if the FlashInfer import fails, serving fails identically at mixer + import (mamba2_mixer.py:23 @1.3.0rc29). + """ + if _serving_selects_flashinfer_ssu(nheads, n_groups, head_dim, d_state): + from flashinfer.mamba import selective_state_update as selective_state_update_flashinfer + + return selective_state_update_flashinfer, SSU_KERNEL_SOURCE_FLASHINFER + return selective_state_update, SSU_KERNEL_SOURCE_TRITON + + +def _decode_dt(batch_size: int, nheads: int, head_dim: int, device) -> torch.Tensor: + """Decode dt laid out the way serving hands it to the state-update kernel. + + Serving casts the per-head dt to fp32 and broadcasts it over head_dim with + einops.repeat (mamba2_mixer.py:743-746 and :764 @1.3.0rc29): an expanded + VIEW with strides (nheads, 1, 0), like the A / D / dt_bias buffers. FlashInfer's + binding checks exactly that layout (csrc/selective_state_update.cu:187-188 + @v0.6.18: dt.stride(1) == 1, dt.stride(2) == 0 — the b200 diagnostic on + 76ce215e failed every decode cell with "dt.stride(1) must be 1, got 64"), and + TRT-LLM's Triton kernel keys its tie_hdim fast path off dt.stride(-1) == 0. + A dense (batch, nheads, head_dim) tensor is rejected by the former and + mis-tuned for the latter. + """ + return repeat(torch.randn(batch_size, nheads, dtype=torch.float32, device=device), "b h -> b h p", p=head_dim) + + def run_mamba2_context_benchmark( d_model: int, d_state: int, @@ -146,7 +257,9 @@ def run_mamba2_context_benchmark( This benchmarks: 1. causal_conv1d_fn - Conv1D for context phase - 2. mamba_chunk_scan_combined - SSM scan + 2. mamba_chunk_scan_combined - SSM scan (the one entry point serving calls; + it routes to FlashInfer's SSD on SM100-family parts itself, + ssd_combined.py:469-497 @1.3.0rc29) Together these represent the core compute of a Mamba2 layer (excluding in_proj/out_proj GEMMs which use existing GEMM model). @@ -176,6 +289,17 @@ def run_mamba2_context_benchmark( D = torch.randn(nheads, device=device) # noqa: N806 dt_bias = torch.rand(nheads, device=device) - 4.0 + # Every cell is timed under CUDA-graph replay; a failed capture raises + # (benchmark_with_power default allow_graph_fail=False) and the cell is + # counted as failed below. Until 2026-10-12 allow_graph_fail=True timed the + # Python launch loop eagerly and shipped that floor as the kernel latency, + # and the per-cell `except` dropped a failed kernel lane with no failure + # record: every shipped mamba2 table (rc20, all eight systems) and the rc29 + # b200 package carry only the causal_conv1d lanes. + successful_points = 0 + failed_points = 0 + failures: list[str] = [] + # Sweep over batch_size x seq_len combinations for batch_size in batch_size_list: for seq_len in seq_len_list: @@ -215,20 +339,34 @@ def run_mamba2_context_benchmark( C = torch.randn(batch_size, seq_len, n_groups, d_state, dtype=dtype, device=device) # noqa: N806 # --- Benchmark causal_conv1d_fn --- + conv_kernel_source = "causal_conv1d_fn" + if _token_major_prefill_conv(): + conv_kernel_source = "causal_conv1d_fn_token_major" + x_tm, out_tm, cu_seqlens, cache_idx, has_init = _prefill_conv_token_major_inputs( + xbc_input, batch_size, seq_len, device + ) + del xbc_input + + def run_conv1d( + _x=x_tm, _o=out_tm, _conv_state=conv_state, _cu=cu_seqlens, _ci=cache_idx, _hi=has_init + ): + _run_prefill_conv_token_major(_x, _o, conv_weight, conv_bias, _conv_state, _cu, _ci, _hi) + + else: + + def run_conv1d(_xbc=xbc_input, _conv_state=conv_state): + causal_conv1d_fn(_xbc, conv_weight, conv_bias, activation="silu", conv_states=_conv_state) + torch.cuda.synchronize() - causal_conv1d_fn(xbc_input, conv_weight, conv_bias, activation="silu", conv_states=conv_state) + run_conv1d() torch.cuda.synchronize() - def run_conv1d(_xbc=xbc_input, _conv_state=conv_state): - causal_conv1d_fn(_xbc, conv_weight, conv_bias, activation="silu", conv_states=_conv_state) - with benchmark_with_power( device=device, kernel_func=run_conv1d, num_warmups=num_warmups, num_runs=num_runs, repeat_n=1, - allow_graph_fail=True, ) as results: log_perf( item_list=[{**common_log_data, "latency": results["latency_ms"]}], @@ -236,12 +374,15 @@ def run_conv1d(_xbc=xbc_input, _conv_state=conv_state): version=tensorrt_llm.__version__, device_name=torch.cuda.get_device_name(device), op_name="mamba2", - kernel_source="causal_conv1d_fn", + kernel_source=conv_kernel_source, perf_filename=perf_filename, power_stats=results["power_stats"], ) # --- Benchmark mamba_chunk_scan_combined --- + # Caller-owned output, as serving's preallocated buffer view + # (mamba2_mixer.py:634 @1.3.0rc29); the kernel asserts on it. + ssm_out = torch.empty_like(x) torch.cuda.synchronize() mamba_chunk_scan_combined( x, @@ -255,10 +396,11 @@ def run_conv1d(_xbc=xbc_input, _conv_state=conv_state): dt_bias=dt_bias, dt_softplus=True, return_final_states=True, + out=ssm_out, ) torch.cuda.synchronize() - def run_ssm_scan(_x=x, _dt=dt, _b=B, _c=C): + def run_ssm_scan(_x=x, _dt=dt, _b=B, _c=C, _out=ssm_out): mamba_chunk_scan_combined( _x, _dt, @@ -271,6 +413,7 @@ def run_ssm_scan(_x=x, _dt=dt, _b=B, _c=C): dt_bias=dt_bias, dt_softplus=True, return_final_states=True, + out=_out, ) with benchmark_with_power( @@ -279,7 +422,6 @@ def run_ssm_scan(_x=x, _dt=dt, _b=B, _c=C): num_warmups=num_warmups, num_runs=num_runs, repeat_n=1, - allow_graph_fail=True, ) as results: log_perf( item_list=[{**common_log_data, "latency": results["latency_ms"]}], @@ -307,29 +449,55 @@ def run_ssm_scan(_x=x, _dt=dt, _b=B, _c=C): device, ) + conv_kernel_source = "causal_conv1d_fn" + conv1d_iter_idx = [0] + if _token_major_prefill_conv(): + conv_kernel_source = "causal_conv1d_fn_token_major" + tm_pool = [] + for _xb in input_pool["xbc"]: + x_tm, _, cu_seqlens, cache_idx, has_init = _prefill_conv_token_major_inputs( + _xb, batch_size, seq_len, device + ) + tm_pool.append(x_tm) + input_pool["xbc"] = tm_pool # channel-major originals released + out_tm = torch.empty_like(tm_pool[0]) + + def run_conv1d( + _pool=tm_pool, + _o=out_tm, + _conv_state=conv_state, + _idx=conv1d_iter_idx, + _cu=cu_seqlens, + _ci=cache_idx, + _hi=has_init, + ): + idx = _idx[0] % total_iters + _idx[0] += 1 + _run_prefill_conv_token_major( + _pool[idx], _o, conv_weight, conv_bias, _conv_state, _cu, _ci, _hi + ) + + else: + + def run_conv1d(_pool=input_pool, _conv_state=conv_state, _idx=conv1d_iter_idx): + idx = _idx[0] % total_iters + _idx[0] += 1 + causal_conv1d_fn( + _pool["xbc"][idx], conv_weight, conv_bias, activation="silu", conv_states=_conv_state + ) + # Warmup with first set of inputs torch.cuda.synchronize() - causal_conv1d_fn( - input_pool["xbc"][0], conv_weight, conv_bias, activation="silu", conv_states=conv_state - ) + run_conv1d() + conv1d_iter_idx[0] = 0 torch.cuda.synchronize() - conv1d_iter_idx = [0] - - def run_conv1d(_pool=input_pool, _conv_state=conv_state, _idx=conv1d_iter_idx): - idx = _idx[0] % total_iters - _idx[0] += 1 - causal_conv1d_fn( - _pool["xbc"][idx], conv_weight, conv_bias, activation="silu", conv_states=_conv_state - ) - with benchmark_with_power( device=device, kernel_func=run_conv1d, num_warmups=num_warmups, num_runs=num_runs, repeat_n=1, - allow_graph_fail=True, ) as results: log_perf( item_list=[{**common_log_data, "latency": results["latency_ms"]}], @@ -337,12 +505,15 @@ def run_conv1d(_pool=input_pool, _conv_state=conv_state, _idx=conv1d_iter_idx): version=tensorrt_llm.__version__, device_name=torch.cuda.get_device_name(device), op_name="mamba2", - kernel_source="causal_conv1d_fn", + kernel_source=conv_kernel_source, perf_filename=perf_filename, power_stats=results["power_stats"], ) # --- Benchmark mamba_chunk_scan_combined --- + # One caller-owned output buffer for every iteration, as + # serving's preallocated view (mamba2_mixer.py:634 @1.3.0rc29). + ssm_out = torch.empty_like(input_pool["x"][0]) torch.cuda.synchronize() mamba_chunk_scan_combined( input_pool["x"][0], @@ -356,12 +527,13 @@ def run_conv1d(_pool=input_pool, _conv_state=conv_state, _idx=conv1d_iter_idx): dt_bias=dt_bias, dt_softplus=True, return_final_states=True, + out=ssm_out, ) torch.cuda.synchronize() ssm_iter_idx = [0] - def run_ssm_scan(_pool=input_pool, _idx=ssm_iter_idx): + def run_ssm_scan(_pool=input_pool, _idx=ssm_iter_idx, _out=ssm_out): idx = _idx[0] % total_iters _idx[0] += 1 @@ -377,6 +549,7 @@ def run_ssm_scan(_pool=input_pool, _idx=ssm_iter_idx): dt_bias=dt_bias, dt_softplus=True, return_final_states=True, + out=_out, ) with benchmark_with_power( @@ -385,7 +558,6 @@ def run_ssm_scan(_pool=input_pool, _idx=ssm_iter_idx): num_warmups=num_warmups, num_runs=num_runs, repeat_n=1, - allow_graph_fail=True, ) as results: log_perf( item_list=[{**common_log_data, "latency": results["latency_ms"]}], @@ -400,16 +572,27 @@ def run_ssm_scan(_pool=input_pool, _idx=ssm_iter_idx): # Cleanup if aic_cached_inputs: - del x, dt, B, C, xbc_input, conv_state + del x, dt, B, C, xbc_input, conv_state, ssm_out else: - del input_pool, conv_state + del input_pool, conv_state, ssm_out gc.collect() torch.cuda.empty_cache() + successful_points += 1 except Exception as e: + failed_points += 1 + failures.append(f"batch_size={batch_size}, seq_len={seq_len}: {type(e).__name__}: {e}") print(f" Error at batch_size={batch_size}, seq_len={seq_len}: {e}") continue + summary = f"ok={successful_points} error={failed_points} skip=0" + print(f"Mamba2 context summary: {summary}") + if failed_points or successful_points == 0: + raise RuntimeError( + f"TRTLLM Mamba2 context collection failed strict completeness: {summary}; " + f"failed cells: {'; '.join(failures)}" + ) + def run_mamba2_generation_benchmark( d_model: int, @@ -429,7 +612,9 @@ def run_mamba2_generation_benchmark( This benchmarks: 1. causal_conv1d_update - Conv1D state update for decode - 2. selective_state_update - SSM state update + 2. the SSM state update serving selects for this geometry: FlashInfer's + selective_state_update or TRT-LLM's Triton one (see + _decode_state_update_kernel); kernel_source names the one that ran Together these represent the core compute of a Mamba2 layer during decode (excluding in_proj/out_proj GEMMs which use existing GEMM model). @@ -464,6 +649,14 @@ def run_mamba2_generation_benchmark( D = repeat(d_base, "h -> h p", p=head_dim) # noqa: N806 dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + # The decode kernel serving's gate selects for this geometry, and its label. + ssu_func, ssu_kernel_source = _decode_state_update_kernel(nheads, n_groups, head_dim, d_state) + + # Same contract as the context runner: graph-timed or counted as failed. + successful_points = 0 + failed_points = 0 + failures: list[str] = [] + for batch_size in batch_size_list: if aic_debug == 1: print(f" Benchmarking batch_size={batch_size}") @@ -477,6 +670,11 @@ def run_mamba2_generation_benchmark( # SSM state: [batch, nheads, head_dim, d_state] ssm_state = torch.randn(batch_size, nheads, head_dim, d_state, dtype=dtype, device=device) + # Serving addresses the state cache through slot indices and hands + # the kernel a view of its preallocated output buffer + # (mamba2_mixer.py:901-907 @1.3.0rc29); the state here is dense, so + # the slots are the identity. + state_batch_indices = torch.arange(batch_size, dtype=torch.int32, device=device) if aic_cached_inputs: # Common log data for both kernels @@ -497,9 +695,10 @@ def run_mamba2_generation_benchmark( # Cached mode: use same inputs for all iterations (saves memory) xbc_input = torch.randn(batch_size, conv_dim, dtype=dtype, device=device) x = torch.randn(batch_size, nheads, head_dim, dtype=dtype, device=device) - dt = torch.randn(batch_size, nheads, head_dim, dtype=dtype, device=device) + dt = _decode_dt(batch_size, nheads, head_dim, device) B = torch.randn(batch_size, n_groups, d_state, dtype=dtype, device=device) # noqa: N806 C = torch.randn(batch_size, n_groups, d_state, dtype=dtype, device=device) # noqa: N806 + ssu_out = torch.empty_like(x) # --- Benchmark causal_conv1d_update --- torch.cuda.synchronize() @@ -515,7 +714,6 @@ def run_conv1d_update(_xbc=xbc_input, _conv_state=conv_state): num_warmups=num_warmups, num_runs=num_runs, repeat_n=1, - allow_graph_fail=True, ) as results: log_perf( item_list=[{**common_log_data, "latency": results["latency_ms"]}], @@ -528,9 +726,9 @@ def run_conv1d_update(_xbc=xbc_input, _conv_state=conv_state): power_stats=results["power_stats"], ) - # --- Benchmark selective_state_update --- + # --- Benchmark the serving-selected selective_state_update --- torch.cuda.synchronize() - selective_state_update( + ssu_func( ssm_state, x, dt, @@ -541,11 +739,13 @@ def run_conv1d_update(_xbc=xbc_input, _conv_state=conv_state): z=None, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=state_batch_indices, + out=ssu_out, ) torch.cuda.synchronize() - def run_state_update(_ssm_state=ssm_state, _x=x, _dt=dt, _b=B, _c=C): - selective_state_update( + def run_state_update(_ssm_state=ssm_state, _x=x, _dt=dt, _b=B, _c=C, _out=ssu_out): + ssu_func( _ssm_state, _x, _dt, @@ -556,6 +756,8 @@ def run_state_update(_ssm_state=ssm_state, _x=x, _dt=dt, _b=B, _c=C): z=None, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=state_batch_indices, + out=_out, ) with benchmark_with_power( @@ -564,7 +766,6 @@ def run_state_update(_ssm_state=ssm_state, _x=x, _dt=dt, _b=B, _c=C): num_warmups=num_warmups, num_runs=num_runs, repeat_n=1, - allow_graph_fail=True, ) as results: log_perf( item_list=[{**common_log_data, "latency": results["latency_ms"]}], @@ -572,7 +773,7 @@ def run_state_update(_ssm_state=ssm_state, _x=x, _dt=dt, _b=B, _c=C): version=tensorrt_llm.__version__, device_name=torch.cuda.get_device_name(device), op_name="mamba2", - kernel_source="selective_state_update", + kernel_source=ssu_kernel_source, perf_filename=perf_filename, power_stats=results["power_stats"], ) @@ -597,7 +798,6 @@ def run_state_update(_ssm_state=ssm_state, _x=x, _dt=dt, _b=B, _c=C): { "xbc": (batch_size, conv_dim), "x": (batch_size, nheads, head_dim), - "dt": (batch_size, nheads, head_dim), "B": (batch_size, n_groups, d_state), "C": (batch_size, n_groups, d_state), }, @@ -605,6 +805,8 @@ def run_state_update(_ssm_state=ssm_state, _x=x, _dt=dt, _b=B, _c=C): dtype, device, ) + input_pool["dt"] = [_decode_dt(batch_size, nheads, head_dim, device) for _ in range(total_iters)] + ssu_out = torch.empty_like(input_pool["x"][0]) # --- Benchmark causal_conv1d_update --- torch.cuda.synchronize() @@ -624,7 +826,6 @@ def run_conv1d_update(_pool=input_pool, _conv_state=conv_state, _idx=conv1d_iter num_warmups=num_warmups, num_runs=num_runs, repeat_n=1, - allow_graph_fail=True, ) as results: log_perf( item_list=[{**common_log_data, "latency": results["latency_ms"]}], @@ -637,9 +838,9 @@ def run_conv1d_update(_pool=input_pool, _conv_state=conv_state, _idx=conv1d_iter power_stats=results["power_stats"], ) - # --- Benchmark selective_state_update --- + # --- Benchmark the serving-selected selective_state_update --- torch.cuda.synchronize() - selective_state_update( + ssu_func( ssm_state, input_pool["x"][0], input_pool["dt"][0], @@ -650,15 +851,17 @@ def run_conv1d_update(_pool=input_pool, _conv_state=conv_state, _idx=conv1d_iter z=None, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=state_batch_indices, + out=ssu_out, ) torch.cuda.synchronize() ssm_iter_idx = [0] - def run_state_update(_pool=input_pool, _ssm_state=ssm_state, _idx=ssm_iter_idx): + def run_state_update(_pool=input_pool, _ssm_state=ssm_state, _idx=ssm_iter_idx, _out=ssu_out): idx = _idx[0] % total_iters _idx[0] += 1 - selective_state_update( + ssu_func( _ssm_state, _pool["x"][idx], _pool["dt"][idx], @@ -669,6 +872,8 @@ def run_state_update(_pool=input_pool, _ssm_state=ssm_state, _idx=ssm_iter_idx): z=None, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=state_batch_indices, + out=_out, ) with benchmark_with_power( @@ -677,7 +882,6 @@ def run_state_update(_pool=input_pool, _ssm_state=ssm_state, _idx=ssm_iter_idx): num_warmups=num_warmups, num_runs=num_runs, repeat_n=1, - allow_graph_fail=True, ) as results: log_perf( item_list=[{**common_log_data, "latency": results["latency_ms"]}], @@ -685,23 +889,34 @@ def run_state_update(_pool=input_pool, _ssm_state=ssm_state, _idx=ssm_iter_idx): version=tensorrt_llm.__version__, device_name=torch.cuda.get_device_name(device), op_name="mamba2", - kernel_source="selective_state_update", + kernel_source=ssu_kernel_source, perf_filename=perf_filename, power_stats=results["power_stats"], ) # Cleanup if aic_cached_inputs: - del ssm_state, conv_state, x, dt, B, C, xbc_input + del ssm_state, conv_state, x, dt, B, C, xbc_input, ssu_out, state_batch_indices else: - del ssm_state, conv_state, input_pool + del ssm_state, conv_state, input_pool, ssu_out, state_batch_indices gc.collect() torch.cuda.empty_cache() + successful_points += 1 except Exception as e: + failed_points += 1 + failures.append(f"batch_size={batch_size}: {type(e).__name__}: {e}") print(f" Error at batch_size={batch_size}: {e}") continue + summary = f"ok={successful_points} error={failed_points} skip=0" + print(f"Mamba2 generation summary: {summary}") + if failed_points or successful_points == 0: + raise RuntimeError( + f"TRTLLM Mamba2 generation collection failed strict completeness: {summary}; " + f"failed cells: {'; '.join(failures)}" + ) + def _make_input_pool( shapes: dict[str, tuple[int, ...]], diff --git a/python/aisimulate/collector/trtllm/collect_mhc_module.py b/python/aisimulate/collector/trtllm/collect_mhc_module.py index 26dc00dcec..0f0a7e9765 100644 --- a/python/aisimulate/collector/trtllm/collect_mhc_module.py +++ b/python/aisimulate/collector/trtllm/collect_mhc_module.py @@ -45,14 +45,14 @@ of the approximation if fused_hc were free, NOT a measured Blackwell fused/unfused delta (which is unmeasured; owner-accepted). Direction is conservative (cost-overestimating). Governed standards exception: -collector/evidence_exceptions.yaml standards_exceptions (expires -2026-10-31); retirement = the coordinated fused_hc table + SDK consumer +collector/evidence_exceptions.yaml standards_exceptions (verified on +1.3.0rc29); retirement = the coordinated fused_hc table + SDK consumer switch for SM100/103. """ from __future__ import annotations -__compat__ = "trtllm>=1.3.0rc20" +__compat__ = "trtllm==1.3.0rc29" import argparse import os diff --git a/python/aisimulate/collector/trtllm/collect_mla.py b/python/aisimulate/collector/trtllm/collect_mla.py index a7faa6105f..28b1d8b93e 100644 --- a/python/aisimulate/collector/trtllm/collect_mla.py +++ b/python/aisimulate/collector/trtllm/collect_mla.py @@ -8,7 +8,7 @@ metadata objects, request/cache helpers, and backend-specific timing paths. """ -__compat__ = "trtllm>=1.3.0rc20" +__compat__ = "trtllm==1.3.0rc29" import math from dataclasses import dataclass @@ -17,13 +17,29 @@ import torch from collector.case_generator import get_context_mla_case_specs, get_generation_mla_case_specs from collector.helper import benchmark_with_power, get_sm_version, log_perf -from tensorrt_llm._torch.attention_backend.interface import ( - AttentionInputType, - MLAParams, - PositionalEmbeddingParams, - RopeParams, -) -from tensorrt_llm._torch.attention_backend.utils import get_attention_backend + +# trtllm >=1.3.0rc29 moved _torch.attention_backend.* -> _torch.attention.backends.* and +# _torch.modules.fused_moe -> _torch.moe.fused_moe (the old package root is a deprecation shim +# without submodules). Path-only compat: same classes, same kernels (layer_permissions.md +# 'API-compat shims may only change HOW the same kernel is constructed'). +try: + from tensorrt_llm._torch.attention.backends.interface import ( + AttentionInputType, + MLAParams, + PositionalEmbeddingParams, + RopeParams, + ) +except ModuleNotFoundError: # < rc29 layout + from tensorrt_llm._torch.attention_backend.interface import ( + AttentionInputType, + MLAParams, + PositionalEmbeddingParams, + RopeParams, + ) +try: + from tensorrt_llm._torch.attention.backends.utils import get_attention_backend +except ModuleNotFoundError: # < rc29 layout + from tensorrt_llm._torch.attention_backend.utils import get_attention_backend from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.bindings.executor import KvCacheConfig @@ -208,14 +224,14 @@ def run_mla( # (cpp/kernels/fmha_v2/setup.py:6926-6931@v1.3.0rc20); the hand-added # sm89 576x512 generation cubin entry from 1.0.0 was dropped in the same # window. Layout-level, so it kills every dtype at once. Fail closed with - # a cited, classified raise (Gemma4/DSA precedent). Re-verify on the next - # framework version bump. + # a cited, classified raise (Gemma4/DSA precedent). Re-verified on 1.3.0rc29 / L40 2026-10-04 (the module + # collector's twin guard, collect_mla_module.py): same assert, now attentionOp.cpp:3234 — still SM89-less. if get_sm_version() < 90: phase = "context" if is_context_phase else "generation" raise ValueError( - f"TRT-LLM MLA has no pre-Hopper FMHA kernel; DeepSeek MLA {phase} " + f"FIXME(kernel-limit): TRT-LLM MLA has no pre-Hopper FMHA kernel; DeepSeek MLA {phase} " f"is unsupported on SM{get_sm_version()} (attentionOp.cpp:" - f"{'3097' if is_context_phase else '3091'} assert @1.3.0rc20)" + f"{'3240' if is_context_phase else '3234'} assert @1.3.0rc29)" ) scenario = Scenario() q_lora_rank = scenario.q_lora_rank @@ -247,6 +263,15 @@ def run_mla( assert num_heads % tp_size == 0, "num_heads != N * tp_size" num_heads = num_heads // tp_size num_kv_heads = num_heads + # FIXME(kernel-limit): TRT-LLM's MLA decode FMHA on SM100/103 tiles 16 q-heads per CTA and + # rejects head counts above 16 that are not a multiple of 16 ("Internal error numHeadsQ=24, + # numHeadsPerCta=16", 724 cases = every heads-24 decode cell of the 96-head model at tp4, + # b200 rc29 job 469988228). Classified before the layer is built. + if not is_context_phase and get_sm_version() in (100, 103) and num_heads > 16 and num_heads % 16: + raise RuntimeError( + "FIXME(kernel-limit): TRT-LLM MLA decode on SM100/103 needs per-rank heads <= 16 or a multiple " + f"of 16 (numHeadsPerCta=16); got {num_heads}" + ) context_sequence_lengths = [input_len for _ in range(batch_size)] diff --git a/python/aisimulate/collector/trtllm/collect_mla_bmm.py b/python/aisimulate/collector/trtllm/collect_mla_bmm.py index 272bcced0b..8d970c5dde 100644 --- a/python/aisimulate/collector/trtllm/collect_mla_bmm.py +++ b/python/aisimulate/collector/trtllm/collect_mla_bmm.py @@ -31,14 +31,13 @@ and logs MLA BMM perf rows for pre/post generation operations. """ -__compat__ = "trtllm>=1.3.0rc20" +__compat__ = "trtllm==1.3.0rc29" import tensorrt_llm import tensorrt_llm.quantization.utils.fp8_utils as fp8_utils import torch -from case_generator import get_mla_bmm_case_specs - -from helper import benchmark_with_power, get_sm_version, log_perf +from collector.case_generator import get_mla_bmm_case_specs +from collector.helper import benchmark_with_power, get_sm_version, log_perf def _supported_dtypes() -> set[str]: diff --git a/python/aisimulate/collector/trtllm/collect_mla_module.py b/python/aisimulate/collector/trtllm/collect_mla_module.py index fac6ba25a3..4c13838aaa 100644 --- a/python/aisimulate/collector/trtllm/collect_mla_module.py +++ b/python/aisimulate/collector/trtllm/collect_mla_module.py @@ -40,12 +40,13 @@ # hardware-observed on SM103 (2026-07-19, B300): identical signature # (SmemTile.cpp:53 "Num. rows must be a multiple of 8 ") for h in # {1,2,4} while h=8 ctx+gen pass — the limit covers the trtllm-gen path on -# both SM100 and SM103. Cases stay as classified runtime failures. +# both SM100 and SM103 (StoreSmemP.h:125 @1.3.0rc29: NumRows >= 8). +# run_mla_module_worker raises for them before the module is built. # Re-verify on the next version bump. # 1.3.0rc20 renamed the cache-manager sparse kwargs and moved attention # metadata to lowered SparseMetadataParams; this module follows those APIs. -__compat__ = "trtllm>=1.3.0rc20" +__compat__ = "trtllm==1.3.0rc29" """ MLA Module Collector for TRT-LLM — unified MLA and DSA benchmarking. @@ -82,13 +83,25 @@ import gc import os import sys +import time import traceback import weakref import tensorrt_llm import torch -from tensorrt_llm._torch.attention_backend.interface import AttentionRuntimeFeatures -from tensorrt_llm._torch.attention_backend.utils import get_attention_backend + +# trtllm >=1.3.0rc29 moved _torch.attention_backend.* -> _torch.attention.backends.* and +# _torch.modules.fused_moe -> _torch.moe.fused_moe (the old package root is a deprecation shim +# without submodules). Path-only compat: same classes, same kernels (layer_permissions.md +# 'API-compat shims may only change HOW the same kernel is constructed'). +try: + from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures +except ModuleNotFoundError: # < rc29 layout + from tensorrt_llm._torch.attention_backend.interface import AttentionRuntimeFeatures +try: + from tensorrt_llm._torch.attention.backends.utils import get_attention_backend +except ModuleNotFoundError: # < rc29 layout + from tensorrt_llm._torch.attention_backend.utils import get_attention_backend from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_deepseekv3 import ( @@ -125,7 +138,13 @@ sys.path.insert(0, _REPO_ROOT) from collector.case_generator import get_mla_module_model_specs, get_mla_module_sweep_spec -from collector.helper import _resolve_local_model_path, benchmark_with_power, get_sm_version, log_perf +from collector.helper import ( + PeriodicGarbageCollector, + _resolve_local_model_path, + benchmark_with_power, + get_sm_version, + log_perf, +) from collector.registry_types import PerfFile if "glm_moe_dsa" not in _CONFIG_REGISTRY: @@ -259,13 +278,31 @@ def _build_module_test_cases(attn_type: str, mode: str): base_cases = get_context_test_cases(attn_type) if mode == "context" else get_generation_test_cases(attn_type) model_specs = get_mla_module_model_specs(attention_type=attn_type, backend="trtllm") cases = [] + dropped_non_shard_heads = 0 for model_spec in model_specs: + # Model-correlated narrowing (case_authoring.md; get_msa_head_counts + # precedent): a row declaring its real tensor_parallel_sizes measures + # only the per-rank head counts those shards produce; the shared head + # grid otherwise sweeps 128..1 for the 64-head GLM-5.x too. Rows without + # the declaration keep the full grid. + shard_heads = None + declared_tp_sizes = getattr(model_spec, "tensor_parallel_sizes", None) + if declared_tp_sizes: + shard_heads = {model_spec.native_num_heads // tp for tp in declared_tp_sizes} for base_case in base_cases: s, b, h, kv_dtype, compute_dtype, gemm_type, *rest = base_case + if shard_heads is not None and h not in shard_heads: + dropped_non_shard_heads += 1 + continue case = [s, b, h, kv_dtype, compute_dtype, gemm_type, model_spec.model_path, attn_type] if rest: case.append(rest[0]) cases.append(case) + if dropped_non_shard_heads: + print( + f"[trtllm-mla-module-cases] {attn_type}/{mode}: dropped {dropped_non_shard_heads} cases whose " + "per-rank heads no declared tensor_parallel_sizes shard produces" + ) return cases @@ -368,6 +405,12 @@ def _apply_gemm_type_quant(model_config, gemm_type: str, use_fp8_kv_cache: bool) _replace_quant_config( model_config.quant_config, quant_algo=QuantAlgo.NVFP4, + # NVFP4 scale blocks are 16 elements; without an explicit value the replaced + # config keeps the checkpoint's/fp8_block group_size=128 and 1.3.0rc29 + # NVFP4LinearMethod.resolve_scaling_vector_size rejects module construction + # ("supports NVFP4 scale blocks of (16,) elements, but the checkpoint declares + # group_size=128" — B200 smoke 2026-10-04; nvfp4 is a platform floor on Hopper). + group_size=16, kv_cache_quant_algo=kv_algo, exclude_modules=None, ), @@ -612,6 +655,39 @@ def create_attention_layer( # ═══════════════════════════════════════════════════════════════════════ +# One attention layer per (model, heads, kv dtype, gemm type) per worker. The layer +# build (ModelConfig.from_pretrained, DeepseekV3DecoderLayer, dummy weights, +# post_load_weights) does not depend on the point: create_attention_layer sizes +# the config with fixed max_num_tokens / max_seq_len, so one layer serves every +# (batch, seq, prefix) of its key exactly as serving reuses one module. The +# per-case KV cache manager, metadata, dry run, CUDA-graph capture and timing are +# untouched. Same pattern as collect_dsv4_attn._cached_dsv4_attention_module +# (owner-approved perf change 2026-08-09); b300 base run (job 471159002): +# dsa_context_module 1.2-2.4 GPU-s and dsa_generation_module 1.8 GPU-s per +# one-row case, most of it this build. At most one layer is retained. +_LAYER_CACHE: dict = {} + + +def _cached_attention_layer(model_path: str, num_heads: int, use_fp8_kv_cache: bool, gemm_type: str, device: str): + key = (model_path, int(num_heads), bool(use_fp8_kv_cache), gemm_type, device) + hit = _LAYER_CACHE.get(key) + if hit is not None: + return hit + if _LAYER_CACHE: + _LAYER_CACHE.clear() + gc.collect() + torch.cuda.empty_cache() + entry = create_attention_layer( + model_path=model_path, + num_heads=num_heads, + use_fp8_kv_cache=use_fp8_kv_cache, + gemm_type=gemm_type, + device=device, + ) + _LAYER_CACHE[key] = entry + return entry + + def create_kv_cache_and_metadata( model_config: ModelConfig, attn_module, @@ -812,6 +888,13 @@ def run_mla_module( # (Gemma4 dense-attention precedent in collect_attn.py) instead of # paying construction + a CUDA-level abort per case. Re-verify against # deepgemm's supported-arch set on the next framework version bump. + # FIXME(re-verify): NOT re-verified at 1.3.0rc29, and the source no longer + # supports it. rc29 pins DeepGEMM 8b1392b (3rdparty/fetch_content.json), + # whose attention.hpp dispatches SM120 for both indexer phases + # (sm120_mqa_logits :205-212, sm120_paged_mqa_logits :414-418; heads + # 16/32/64). Until an SM120 rc29 run decides it, the raise stays and + # keeps its unprefixed message, so the failure log reports it as + # unexpected instead of an accepted kernel limit. if attn_type == "dsa" and get_sm_version() >= 120: raise ValueError( f"DeepGEMM sparse-attention indexer has no RTX Blackwell kernel; DSA " @@ -823,12 +906,15 @@ def run_mla_module( # (attentionOp.cpp:3097) / generation (:3091). Hardware-observed on L40 # (SM89) 2026-07-26: module smoke 0/8, C++ SIGABRT per case; serving hits # the identical assert. Same fail-closed treatment as the stock MLA - # collector (collect_mla.py). Re-verify on the next framework version bump. + # collector (collect_mla.py). Re-verified on 1.3.0rc29 / L40 (SM89) 2026-10-04: with only this SM guard + # patched out, a bf16 DeepSeek-V3 MLA module cell still aborts in AttentionOp with "Deepseek should be + # supported by fmha in generation part" (attentionOp.cpp:3234 @1.3.0rc29; the assert moved from :3091), SIGSEGV + # in the worker (facts/mla_bypass_*.log). The wall is unchanged across rc20 -> rc29: keep failing closed. if attn_type == "mla" and get_sm_version() < 90: raise ValueError( - f"TRT-LLM MLA has no pre-Hopper FMHA kernel; MLA modules are " + f"FIXME(kernel-limit): TRT-LLM MLA has no pre-Hopper FMHA kernel; MLA modules are " f"unsupported on SM{get_sm_version()} " - f"(attentionOp.cpp:3091/:3097 assert @1.3.0rc20)" + f"(attentionOp.cpp:3234 assert @1.3.0rc29, :3091/:3097 @1.3.0rc20)" ) # FIXME(kernel-limit): SM120 MLA context takes serving's dense-expand path # (forward_context_default, attention.py:2015-2040@1.3.0rc20 — SM100/103 @@ -854,14 +940,23 @@ def run_mla_module( f"compute={compute_dtype}, kv={kv_cache_dtype}, model={model_path}" ) + # Per-phase wall time of this case, printed after cleanup so the campaign + # log shows where a module case spends its seconds (layer cache hit vs + # KV manager + metadata vs graph capture + timing vs teardown). + _t_case = _t_last = time.perf_counter() + phase_timing: dict[str, float] = {} + + def _mark(name: str) -> None: + nonlocal _t_last + now = time.perf_counter() + phase_timing[name] = now - _t_last + _t_last = now + # 1. Create attention layer (from_pretrained reads config.json) - attn_module, model_config, original_architecture = create_attention_layer( - model_path=model_path, - num_heads=num_heads, - use_fp8_kv_cache=use_fp8_kv_cache, - gemm_type=gemm_type, - device=device, + attn_module, model_config, original_architecture = _cached_attention_layer( + model_path, num_heads, use_fp8_kv_cache, gemm_type, device ) + _mark("layer") # 2. Create KV cache + metadata kv_cache_manager, attn_metadata = create_kv_cache_and_metadata( @@ -875,6 +970,8 @@ def run_mla_module( device=device, ) + _mark("kv_meta") + # 3. Input tensors hidden_size = model_config.pretrained_config.hidden_size if is_context: @@ -902,6 +999,8 @@ def run_mla_module( device=torch_device, ) + _mark("inputs") + # 4. Dry run with model_extra_attrs(model_config.extra_attrs): get_model_extra_attrs()["attention_metadata"] = weakref.ref(attn_metadata) @@ -918,12 +1017,20 @@ def run_mla_module( _cleanup(kv_cache_manager) raise + _mark("dry_run") + # 5. Benchmark import tensorrt_llm._torch.utils as _trtllm_utils _trtllm_utils._model_extra_attrs.attrs = model_config.extra_attrs _trtllm_utils._model_extra_attrs.attrs["attention_metadata"] = weakref.ref(attn_metadata) + # FIXME(kernel-limit): same 32-bit per-sequence offset boundary as + # collect_attn.py (H20 campaign 2026-10-02, rc29): DeepSeek-V3 context cells + # with (batch-1) * input_len * local_heads * 192 * 2 >= 2**31 (tp1: 81920+ + # total tokens) fault with cudaErrorIllegalAddress while the TP-sharded rows + # of the same (batch, seq) pass. Unverified against the kernel source; left + # failing into the classified log. def kernel_func(): attn_module.forward(position_ids, hidden_states, attn_metadata) @@ -931,6 +1038,9 @@ def kernel_func(): device=torch_device, kernel_func=kernel_func, num_warmups=warming_up, + # The case's dry run already executed this forward; three eager passes + # cover lazily initialised paths before capture. + num_eager_warmups=3, num_runs=test_ite, repeat_n=1, allow_graph_fail=False, @@ -938,6 +1048,7 @@ def kernel_func(): pass latency = results["latency_ms"] + _mark("bench") # 6. Log results if is_context: @@ -983,6 +1094,13 @@ def kernel_func(): ) _cleanup(kv_cache_manager) + _mark("cleanup") + print( + f" [{variant} module timing] {phase} b={batch_size}, s={seq_len}, prefix={prefix_len}, " + f"heads={num_heads}, gemm={gemm_type}, kv={kv_cache_dtype}: " + + " ".join(f"{name}={seconds:.2f}s" for name, seconds in phase_timing.items()) + + f" total={time.perf_counter() - _t_case:.2f}s" + ) return latency @@ -1001,6 +1119,13 @@ def run_mla_module_worker( device: str = "cuda:0", ): """Worker-compatible positional wrapper used by collector/collect.py.""" + # FIXME(kernel-limit): the SM100 DSA note at the top of this file — per-rank heads 1/2/4 failed + # 15,408 b200 rc29 cases after a full module load each (job 469988228). Raise before the build. + if attn_type == "dsa" and get_sm_version() in (100, 103) and num_heads < 8: + raise RuntimeError( + "FIXME(kernel-limit): TRT-LLM trtllm-gen sparse-MLA FMHA on SM100/103 needs per-rank heads >= 8 " + f"(num_heads={num_heads}; StoreSmemP.h:125 @1.3.0rc29)" + ) return run_mla_module( seq_len=seq_len, batch_size=batch_size, @@ -1016,11 +1141,15 @@ def run_mla_module_worker( ) +# gc.collect() after every case cost 0.17 s of a ~1 s module case; collect periodically instead. +_CASE_GC = PeriodicGarbageCollector() + + def _cleanup(kv_cache_manager): if kv_cache_manager is not None: kv_cache_manager.shutdown() torch.cuda.empty_cache() - gc.collect() + _CASE_GC.maybe_collect() # ═══════════════════════════════════════════════════════════════════════ diff --git a/python/aisimulate/collector/trtllm/collect_moe.py b/python/aisimulate/collector/trtllm/collect_moe.py index 4021cad9df..8fd38cc2bb 100755 --- a/python/aisimulate/collector/trtllm/collect_moe.py +++ b/python/aisimulate/collector/trtllm/collect_moe.py @@ -9,10 +9,11 @@ kernel-specific filters. """ -__compat__ = "trtllm>=1.3.0rc20" +__compat__ = "trtllm==1.3.0rc29" import gc import glob +import importlib import inspect import json import os @@ -26,7 +27,15 @@ from tensorrt_llm._torch.autotuner import AutoTuner, autotune from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Gate -from tensorrt_llm._torch.modules.fused_moe import RenormalizeMoeRoutingMethod, create_moe + +# trtllm >=1.3.0rc29 moved _torch.attention_backend.* -> _torch.attention.backends.* and +# _torch.modules.fused_moe -> _torch.moe.fused_moe (the old package root is a deprecation shim +# without submodules). Path-only compat: same classes, same kernels (layer_permissions.md +# 'API-compat shims may only change HOW the same kernel is constructed'). +try: + from tensorrt_llm._torch.moe.fused_moe import RenormalizeMoeRoutingMethod, create_moe +except ModuleNotFoundError: # < rc29 layout + from tensorrt_llm._torch.modules.fused_moe import RenormalizeMoeRoutingMethod, create_moe from tensorrt_llm._utils import is_sm_100f from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -63,6 +72,45 @@ _TRTLLM_VERSION = tensorrt_llm.__version__ +_DEEPGEMM_WORKSPACE_POOL_WARNED = False + + +def _clear_deepgemm_moe_workspaces() -> None: + """Drop DeepGEMM's fp8_block MoE workspace pool between token counts. + + DeepGemmFusedMoE sizes its workspaces per num_tokens and parks every size + in the memory_buffer_utils pool (``DeepGemmFusedMoE.buffers``); left alone + over the 27-token-count sweep the pool reached 251 GB on gb300 and the + next 63 GiB request OOMed (2026-10-11: gb300 60 / gb200 108 fp8_block + cases, Kimi-K2 / DeepSeek-V3 geometries). 1.3.0rc29 moved the module from + ``_torch.modules.fused_moe`` to ``_torch.moe.fused_moe``; the previous + single-path import failed silently inside a try/except and the pool was + never cleared. Both layouts are tried here and an unreachable pool is + reported once instead of swallowed. + """ + global _DEEPGEMM_WORKSPACE_POOL_WARNED + cls = None + for module_name in ( + "tensorrt_llm._torch.moe.fused_moe.fused_moe_deepgemm", # >= 1.3.0rc29 + "tensorrt_llm._torch.modules.fused_moe.fused_moe_deepgemm", # < rc29 + ): + try: + cls = importlib.import_module(module_name).DeepGemmFusedMoE + break + except (ImportError, AttributeError): + continue + pool = getattr(getattr(cls, "buffers", None), "buffers", None) + if pool is None: + if not _DEEPGEMM_WORKSPACE_POOL_WARNED: + _DEEPGEMM_WORKSPACE_POOL_WARNED = True + print( + "[trtllm-moe] DeepGemmFusedMoE workspace pool is not reachable on this TRT-LLM build; " + "fp8_block MoE workspaces will accumulate across token counts" + ) + return + pool.clear() + + def _moe_model_behavior(model_name: str) -> str: """Resolve model-name branches that change the synthetic invocation.""" if any(pattern in model_name for pattern in NON_GATED_MOE_MODELS): @@ -207,7 +255,9 @@ def get_moe_test_cases(): min_latency_mode_options = [False] - if moe_type == "nvfp4" and get_sm_version() == 100 and common_moe_testcase.num_experts <= 256: + if moe_type == "nvfp4" and get_sm_version() in (100, 103) and common_moe_testcase.num_experts <= 256: + # SM103 treated as SM100 (owner decision 2026-10-05: "sm103 可以认为等价sm100" — + # same capability major, same trtllm-gen kernel family; re-verify on a B300 box). # FIXME: recent version only supports SM100 for min-latency mode. # current support, DS router only support up to 256 experts. # Renormalize router only support <=128 experts. trtllmgen kernels only @@ -270,7 +320,86 @@ def get_moe_test_cases(): return test_cases -def run_moe_torch( +def _declared_moe_routing(model_name: str): + """The model row's routing declaration (cases/models/*_cases.yaml + model_case_values.moe: sglang_moe_routing_method_type / _num_expert_group / + _topk_group / _routed_scaling_factor). Framework-neutral model facts despite + the sglang_ prefix: TensorRT-LLM reads the same numbers from the checkpoint + config (n_group, topk_group, routed_scaling_factor) when it builds the gate.""" + for case in get_common_moe_test_cases(): + if case.model_name == model_name: + return case + return None + + +def _zero_routing_bias(gate) -> None: + """Serving gates allocate e_score_correction_bias with torch.empty and only + load_weights fills it (DeepseekV3Gate, modeling_deepseekv3.py:884,926-935 + @1.3.0rc29; DeepseekV4Gate and MiniMaxM3Gate likewise). Their top-k ranks + score + bias, and the collector's balanced / power-law logits separate the + intended experts by ~1e-3 after the sigmoid, so whatever that memory holds + would pick the experts instead of the row's distribution. Zero it, as + vllm/collect_moe.py and wideep/trtllm/collect_moe_a2a.py do.""" + with torch.no_grad(): + gate.e_score_correction_bias.zero_() + + +def _serving_moe_gate(declared, hidden_size, num_experts, topk, dtype, moe_backend): + """The gate serving builds for this architecture, with a zeroed bias. + + DeepSeek-V3-family / GLM / Kimi / Nemotron-H: DeepseekV3Gate + (modeling_deepseekv3.py:979-994 @1.3.0rc29). DeepSeek-V4: DeepseekV4Gate, + whose routing runs separately (gate_forward, sqrt-softplus scoring; + modeling_deepseekv4.py:1533-1546, moe/fused_moe/routing.py:524-614). The + collected layer is a scored one: the first n_hash_layers route by token id + (tid2eid) and have no logit distribution to label. MiniMax-M3: + MiniMaxM3Gate, MiniMax2 routing (modeling_minimaxm3.py:641-646, + routing.py:697-760). Serving resolves moe_backend upper-case; DeepseekV3Gate + keeps a bf16 bias only for 'TRTLLM' (modeling_deepseekv3.py:880-883). + """ + n_group = int(declared.sglang_moe_num_expert_group or 1) + topk_group = int(declared.sglang_moe_topk_group or 1) + routed_scaling_factor = float(declared.sglang_moe_routed_scaling_factor or 1.0) + backend = str(moe_backend).upper() + if declared.architecture == "DeepseekV4ForCausalLM": + from tensorrt_llm._torch.models.modeling_deepseekv4 import DeepseekV4Gate + + gate = DeepseekV4Gate( + hidden_size, + num_experts, + top_k=topk, + n_group=n_group, + topk_group=topk_group, + routed_scaling_factor=routed_scaling_factor, + is_hashed=False, + dtype=dtype, + moe_backend=backend, + ) + elif declared.architecture == "MiniMaxM3ForCausalLM": + from tensorrt_llm._torch.models.modeling_minimaxm3 import MiniMaxM3Gate + + gate = MiniMaxM3Gate( + hidden_size=hidden_size, + num_experts=num_experts, + top_k=topk, + routed_scaling_factor=routed_scaling_factor, + ) + else: + gate = DeepseekV3Gate( + hidden_size, + num_experts, + top_k=topk, + n_group=n_group, + topk_group=topk_group, + routed_scaling_factor=routed_scaling_factor, + dtype=dtype, + moe_backend=backend, + ) + _zero_routing_bias(gate) + return gate + + +def _run_moe_case( moe_type, num_tokens_lists, hidden_size, @@ -287,7 +416,7 @@ def run_moe_torch( perf_filename, device="cuda:0", ): - """Run MoE forward passes and log latency/power to perf file (trtllm >= 1.1 collector).""" + """Run one MoE case (its token sweep) and log latency/power rows (trtllm >= 1.1 collector).""" device = torch.device(device) torch.cuda.set_device(device) torch.set_default_device(device) @@ -298,18 +427,33 @@ def run_moe_torch( print("MOE Allocated GDRAM:", torch.cuda.memory_allocated(device.index) / 1024**2, "MB") print("MOE Reserved GDRAM:", torch.cuda.memory_reserved(device) / 1024**2, "MB") # moe type support bfloat16, fp8_qdq, fp8_block, w4a8, nvfp4(not implemented yet) + # ``dtype`` is the MODEL/activation dtype serving hands to create_moe + # (modeling_deepseekv3.py:798 ``dtype=config.torch_dtype`` -> DeepseekV3MoE + # -> create_moe(dtype=dtype) @1.3.0rc29; identical at rc23). Quantization + # travels in quant_config only. Passing float8 here (pre-rc29 habit) is + # rejected by rc29's MoE resolution ("CutlassFusedMoE FP8_BLOCK_SCALES + # requires torch.bfloat16, got torch.float8_e4m3fn") and never matched + # serving, where activations are bf16 and quantized inside the kernel path. dtype = torch.bfloat16 quant_group_size = 128 quant_algo = None + # FIXME(kernel-limit): TRT-LLM >= 1.3.0rc29 MoE resolution has no implementation for W4A16 on + # SM100/103 ("no MoE implementation can serve this layer ... CutlassFusedMoE: quant_unsupported + # (quant_algo=W4A16)", 114/114 int4_wo cases, b200 job 469988228). Serving cannot load the + # W4A16 checkpoint either (benchmarks/evidence/collector/op-probe-2026-10/results/sm100/trtllm-1.3.0rc29.yaml: + # Kimi-K2.5 fail, finding sm100_probe_failures_recheck_2026_10_04 (g)). The lane is unreachable on this pin + # (moe_resolution.py:665 @1.3.0rc29 raises it), not a per-shape failure; refuse before building weights. + if moe_type == "int4_wo" and get_sm_version() in (100, 103): + raise RuntimeError( + "FIXME(kernel-limit): TRT-LLM int4_wo (W4A16) MoE has no implementation on SM100/103 at " + f"{_TRTLLM_VERSION}; serving fails identically (Kimi-K2.5 identity probe)" + ) if moe_type == "fp8_block": quant_algo = QuantAlgo.FP8_BLOCK_SCALES - dtype = torch.float8_e4m3fn elif moe_type == "w4afp8": quant_algo = QuantAlgo.W4A8_AWQ - dtype = torch.float8_e4m3fn elif moe_type == "fp8": quant_algo = QuantAlgo.FP8 - dtype = torch.float8_e4m3fn elif moe_type == "int4_wo": quant_algo = QuantAlgo.W4A16 int4_config = get_moe_quantization_module_config("trtllm", moe_type, model_name=model_name) @@ -392,7 +536,7 @@ def run_moe_torch( # next framework version bump. if moe_type == "fp8_block" and (hidden_size % 128 != 0 or (inter_size // moe_tp_size) % 128 != 0): raise ValueError( - f"fp8_block MoE requires 128-aligned hidden_size and TP-sharded intermediate " + f"FIXME(kernel-limit): fp8_block MoE requires 128-aligned hidden_size and TP-sharded intermediate " f"size (128x128-blocked weight scales; deepgemm layout.hpp:78 on SM90/100/103, " f"Triton block-scale on SM120); got hidden_size={hidden_size}, " f"inter_size={inter_size} / moe_tp={moe_tp_size} = {inter_size // moe_tp_size}" @@ -400,7 +544,7 @@ def run_moe_torch( if moe_type == "w4afp8" and (inter_size // moe_tp_size) % 128 != 0: raise ValueError( - f"w4afp8 MoE requires a 128-aligned TP-sharded intermediate size (grouped-GEMM " + f"FIXME(kernel-limit): w4afp8 MoE requires a 128-aligned TP-sharded intermediate size (grouped-GEMM " f"k alignment); got inter_size={inter_size} / moe_tp={moe_tp_size} = " f"{inter_size // moe_tp_size}" ) @@ -426,7 +570,7 @@ def run_moe_torch( native_mxfp4_padding = moe_type in _MXFP4_MOE_TYPES and 100 <= sm_version < 120 if not native_mxfp4_padding and (inter_size // moe_tp_size) % (256 // _weight_bits) != 0: raise ValueError( - f"TRT-LLM fused MoE requires the TP-sharded intermediate size to be a multiple " + f"FIXME(kernel-limit): TRT-LLM fused MoE requires the TP-sharded intermediate size to be a multiple " f"of 256/weight_bits = {256 // _weight_bits} for {moe_type} (TLLM_CHECK_WITH_INFO " f"weight-layout alignment); got inter_size={inter_size} / moe_tp={moe_tp_size} = " f"{inter_size // moe_tp_size}" @@ -493,7 +637,7 @@ def run_moe_torch( topk_group = 4 routed_scaling_factor = 2.5 - routing_method = DeepseekV3Gate( + min_latency_gate = DeepseekV3Gate( hidden_size, num_experts, top_k=topk, @@ -502,17 +646,37 @@ def run_moe_torch( routed_scaling_factor=routed_scaling_factor, dtype=dtype, moe_backend="TRTLLM", - ).routing_method + ) + _zero_routing_bias(min_latency_gate) + routing_method = min_latency_gate.routing_method router_logits_dtype = torch.float32 else: - # for low latency mode in fp4, experts > 128 is not supported. - # Native GPT-OSS uses BF16 routing weights for TRTLLM, FP32 otherwise: - # modeling_gpt_oss.py:159-162@1.3.0rc14 (93cb6518), :165-168@1.3.0rc20. - if is_gpt_oss: + declared = _declared_moe_routing(model_name) + if declared is not None and declared.sglang_moe_routing_method_type == "DeepSeekV3": + # Serving routes DeepSeek-V3-family / GLM / Nemotron-H MoE through + # DeepseekV3Gate.routing_method (noaux_tc, grouped top-k with the + # e_score_correction_bias; Deepseekv3RoutingImpl is_fused=True -> + # deepseek_v3_topk_kernel) built from the checkpoint's n_group / + # topk_group / routed_scaling_factor and the resolved moe_backend — + # modeling_deepseekv3.py:979-990, modeling_glm.py:273-285, + # modeling_nemotron_h.py:236 @1.3.0rc29. RenormalizeMoeRoutingMethod + # (customMoeRoutingKernel) is the softmax-top-k family (Qwen3-MoE, + # Mixtral, ...), and running it for DeepSeek-V3 diverged the + # moe_fp8block gate on H20 (2026-10-01). The router GEMM that feeds the + # gate (dsv3_router_gemm_op) stays outside this op, as before. + # DeepSeek-V4 and MiniMax-M3 build their own gates (_serving_moe_gate). + routing_method = _serving_moe_gate( + declared, hidden_size, num_experts, topk, dtype, model_config.moe_backend + ).routing_method + router_logits_dtype = torch.float32 + elif is_gpt_oss: + # Native GPT-OSS uses BF16 routing weights for TRTLLM, FP32 otherwise: + # modeling_gpt_oss.py:159-162@1.3.0rc14 (93cb6518), :165-168@1.3.0rc20. routing_method = RenormalizeMoeRoutingMethod( topk, output_dtype=torch.bfloat16 if model_config.moe_backend == "trtllm" else torch.float32 ) else: + # softmax top-k renormalized routing (fused custom routing kernel). routing_method = RenormalizeMoeRoutingMethod(topk) create_moe_kwargs = { @@ -725,12 +889,7 @@ def fp32_to_mxfp4(tensor): del hidden_states_max_tokens, logits_max_tokens if moe_type == "fp8_block": - try: - from tensorrt_llm._torch.modules.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE - - DeepGemmFusedMoE.buffers.buffers.clear() - except (ImportError, AttributeError): - pass + _clear_deepgemm_moe_workspaces() torch.cuda.empty_cache() for num_tokens in num_tokens_lists: @@ -766,29 +925,28 @@ def run_forward_pass(): moe.forward(hidden_states, actual_logits, do_finalize=not min_latency_mode) # noqa: F821 # ═══════════════════════════════════════════════════════════════════════════════ - # Benchmark with automatic power measurement and graph fallback + # Benchmark under CUDA-graph replay with automatic power measurement # ═══════════════════════════════════════════════════════════════════════════════ # Determine base warmups and runs based on distribution mode num_warmups = 1 if distributed == "power_law" else 3 num_runs = 1 if distributed == "power_law" else 6 - # Use benchmark_with_power with graceful graph fallback + # CUDA-graph timing is mandatory: a failed capture raises + # (benchmark_with_power default allow_graph_fail=False) and the case + # fails classified. The moe row has no column that could record an + # eager-timed measurement, and the fallback used until 2026-10-12 + # shipped the Python launch-loop floor as kernel latency. with benchmark_with_power( device=device, kernel_func=run_forward_pass, num_warmups=num_warmups, num_runs=num_runs, repeat_n=1, - allow_graph_fail=True, # Enable graceful fallback to eager execution ) as results: # Calculate per-iteration latency (accounting for internal iterations) latency = results["latency_ms"] / num_iter power_stats = results["power_stats"] - # Log if CUDA graph capture failed (for debugging) - if not results["used_cuda_graph"] and aic_debug == 1: - print(f"CUDA graph capture failed for {num_tokens} tokens, used eager execution fallback") - if moe_type == "fp8_block" and is_sm_100f(sm_version): source = "deepgemm" elif moe_type == "fp8_block" and get_sm_version() == 120: @@ -847,18 +1005,168 @@ def run_forward_pass(): del actual_logits del hidden_states if moe_type == "fp8_block" and num_tokens != max_tokens: - try: - from tensorrt_llm._torch.modules.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE - - DeepGemmFusedMoE.buffers.buffers.clear() - except (ImportError, AttributeError): - pass + _clear_deepgemm_moe_workspaces() torch.cuda.empty_cache() - if os.getenv("TRTLLM_MOE_RESTART_WORKER", "1") != "0": - # Exit the worker process after completing MOE task to ensure complete resource cleanup. - # This forces OS to reclaim all GPU memory, CUDA context, and other resources. - sys.exit(EXIT_CODE_RESTART) + +# ── Worker recycling ───────────────────────────────────────────────────────── +# AIConfigurator's collector (40bc989a9) exited the worker after EVERY case so the +# OS reclaimed all GPU memory. collect.py spawns its workers, so each recycle +# re-imports tensorrt_llm and re-creates the CUDA context (10-15 s); on the b300 +# base run (job 471159002) that was most of a 14-20 GPU-s case of 27 token +# points. The policy is now: +# TRTLLM_MOE_RESTART_EVERY=N recycle after N completed cases on this worker +# (default 16; 1 = the legacy per-case recycle; +# 0 = never recycle on the count) +# TRTLLM_MOE_RESTART_LEAK_MB=M recycle at once when more than M MiB are still +# allocated after the case's own cleanup (default +# 1024; 0 disables the guard) — the residue the +# per-case exit used to hide +# TRTLLM_MOE_RESTART_WORKER=0 legacy switch: never recycle (overrides both) +# A recycle is always announced on one greppable line so a run log shows how +# often, and why, workers were replaced. +_cases_since_worker_start = 0 + + +def _worker_restart_policy() -> tuple[int, int] | None: + """(cases per worker, leak guard MiB), or None when recycling is disabled.""" + if os.getenv("TRTLLM_MOE_RESTART_WORKER", "1") == "0": + return None + every = int(os.getenv("TRTLLM_MOE_RESTART_EVERY", "16")) + leak_mb = int(os.getenv("TRTLLM_MOE_RESTART_LEAK_MB", "1024")) + if every < 0 or leak_mb < 0: + raise ValueError( + f"TRTLLM_MOE_RESTART_EVERY and TRTLLM_MOE_RESTART_LEAK_MB must be >= 0; got {every} and {leak_mb}" + ) + return every, leak_mb + + +def _release_trtllm_memory_buffers() -> int: + """Drop TRT-LLM's global workspace buffer pool between cases; returns MiB released. + + b200 probe 2026-10-29 (job 514855922, TRTLLM_MOE_LEAK_DEBUG=1): everything still + allocated after a case was `BufferBlock` tensors of + tensorrt_llm/_torch/memory_buffer_utils.py — the module-level ``_buffer = Buffers()`` + pool the fused-MoE kernels draw their workspaces from (2.7-86 GiB, the two + largest blocks in the 2:1 shape of w3_w1 / w2). It is a reuse pool, not a leak: + the next case replaces it, so memory did not grow across 117 cases. Serving keeps + it alive for the model's lifetime; between collector cases nothing is live, so + every block (reserved ones included — ``clear_memory_buffers`` keeps those) is + released here and the leak guard measures real residue again. @1.3.0rc29; a + TRT-LLM without this module (older pins) makes this a no-op. + """ + try: + from tensorrt_llm._torch import memory_buffer_utils + except ImportError: + return 0 + pool = memory_buffer_utils.get_memory_buffers() + released = 0 + for blocks in list(pool.buffers.values()): + for block in blocks: + tensor = getattr(block, "buffer", None) + if tensor is not None: + released += tensor.numel() * tensor.element_size() + block.buffer = None + pool.buffers.clear() + return released // 2**20 + + +def _print_live_cuda_tensors(device, top: int = 8) -> None: + """TRTLLM_MOE_LEAK_DEBUG=1: name what still pins CUDA memory after a case. + + b200 probe 2026-10-29 (jobs 514778748/514778867): 2.7-86 GiB stayed allocated after + every DeepSeek-V3 fp8_block case, so the leak guard recycled the worker each time and + the policy bought nothing. Lists the largest live CUDA tensors with the types of the + objects that reference them (gc view; a traceback-free best effort). + """ + sizes = [] + for obj in gc.get_objects(): + try: + if isinstance(obj, torch.Tensor) and obj.is_cuda and obj.device == device: + sizes.append((obj.numel() * obj.element_size(), obj)) + except Exception: # gc.get_objects can hand back half-built objects + continue + sizes.sort(key=lambda item: -item[0]) + total = sum(size for size, _ in sizes) + print(f"trtllm.moe: leak debug: {len(sizes)} live CUDA tensors, {total / 2**20:.0f} MiB; largest:") + for size, tensor in sizes[:top]: + referrers = [type(r).__name__ for r in gc.get_referrers(tensor) if r is not sizes][:4] + print( + f" {size / 2**20:8.0f} MiB {tuple(tensor.shape)} {tensor.dtype} " + f"param={isinstance(tensor, torch.nn.Parameter)} referrers={referrers}" + ) + + +def _recycle_worker_if_due(device) -> None: + """Exit this worker with EXIT_CODE_RESTART when the recycling policy says so. + + Runs after the case's rows are persisted; collect.py marks the task done and + starts a fresh worker on the same device (worker(): SystemExit handling). + """ + global _cases_since_worker_start + policy = _worker_restart_policy() + if policy is None: + return + every, leak_mb = policy + _cases_since_worker_start += 1 + workspace_mb = _release_trtllm_memory_buffers() + gc_collect() + allocated_mb = torch.cuda.memory_allocated(device) / 2**20 + reserved_mb = torch.cuda.memory_reserved(device) / 2**20 + print( + f"trtllm.moe: after case {_cases_since_worker_start} on this worker: " + f"{allocated_mb:.0f} MiB allocated, {reserved_mb:.0f} MiB reserved " + f"(released {workspace_mb} MiB of TRT-LLM workspace buffers)" + ) + if os.getenv("TRTLLM_MOE_LEAK_DEBUG", "0") == "1" and allocated_mb > 256: + _print_live_cuda_tensors(device) + due_count = every > 0 and _cases_since_worker_start >= every + due_leak = leak_mb > 0 and allocated_mb > leak_mb + if not (due_count or due_leak): + return + reason = "leak guard" if due_leak else f"every {every}" + print( + f"trtllm.moe: recycling worker after {_cases_since_worker_start} case(s) " + f"({allocated_mb:.0f} MiB still allocated, {reason})" + ) + sys.exit(EXIT_CODE_RESTART) + + +def run_moe_torch( + moe_type, + num_tokens_lists, + hidden_size, + inter_size, + topk, + num_experts, + moe_tp_size, + moe_ep_size, + min_latency_mode, + model_name, + distributed="power_law", + power_law_alpha=0.0, + *, + perf_filename, + device="cuda:0", +): + """Run one MoE case, then recycle the worker when the policy above is due.""" + _run_moe_case( + moe_type, + num_tokens_lists, + hidden_size, + inter_size, + topk, + num_experts, + moe_tp_size, + moe_ep_size, + min_latency_mode, + model_name, + distributed=distributed, + power_law_alpha=power_law_alpha, + perf_filename=perf_filename, + device=device, + ) + _recycle_worker_if_due(torch.device(device)) if __name__ == "__main__": diff --git a/python/aisimulate/collector/trtllm/collect_msa_module.py b/python/aisimulate/collector/trtllm/collect_msa_module.py index f3421b9c9c..a332f4e282 100644 --- a/python/aisimulate/collector/trtllm/collect_msa_module.py +++ b/python/aisimulate/collector/trtllm/collect_msa_module.py @@ -3,7 +3,7 @@ # MiniMax-M3 support landed in TRT-LLM 1.3.0rc19 (modeling_minimaxm3.py + # attention_backend/sparse/minimax_m3/); this collector follows the rc20 APIs. -__compat__ = "trtllm>=1.3.0rc19" +__compat__ = "trtllm==1.3.0rc29" """ MSA Module Collector for TRT-LLM — MiniMax-M3 sparse-attention benchmarking. @@ -55,8 +55,19 @@ import tensorrt_llm import torch import transformers -from tensorrt_llm._torch.attention_backend.interface import AttentionRuntimeFeatures -from tensorrt_llm._torch.attention_backend.utils import get_attention_backend + +# trtllm >=1.3.0rc29 moved _torch.attention_backend.* -> _torch.attention.backends.* and +# _torch.modules.fused_moe -> _torch.moe.fused_moe (the old package root is a deprecation shim +# without submodules). Path-only compat: same classes, same kernels (layer_permissions.md +# 'API-compat shims may only change HOW the same kernel is constructed'). +try: + from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures +except ModuleNotFoundError: # < rc29 layout + from tensorrt_llm._torch.attention_backend.interface import AttentionRuntimeFeatures +try: + from tensorrt_llm._torch.attention.backends.utils import get_attention_backend +except ModuleNotFoundError: # < rc29 layout + from tensorrt_llm._torch.attention_backend.utils import get_attention_backend from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_minimaxm3 import MiniMaxM3DecoderLayer @@ -104,8 +115,15 @@ get_mla_module_model_specs, get_mla_module_precision_specs, get_mla_module_sweep_spec, + get_msa_head_counts, +) +from collector.helper import ( + PeriodicGarbageCollector, + _resolve_local_model_path, + benchmark_with_power, + get_sm_version, + log_perf, ) -from collector.helper import _resolve_local_model_path, benchmark_with_power, get_sm_version, log_perf from collector.registry_types import PerfFile if not hasattr(_trtllm_configs, "PretrainedConfig"): @@ -134,6 +152,18 @@ def _get_precision_combos(phase: str): ] +def _msa_implementation(sm_version: int) -> str: + """Which MiniMax-M3 sparse implementation TRT-LLM can run on this SM. + + "msa" = the fmha_sm100 kernels, hard-gated to the SM100 family + (msa_availability.py:ensure_msa_available@1.3.0rc29 raises off SM100/103, + SM120 included); everything else runs the Triton reference + (triton_backend.py). Single source of truth for the plan-time budget + below and the layer construction in build_msa_layer. + """ + return "msa" if sm_version in (100, 103) else "triton" + + def get_context_test_cases(): """Context-phase test cases. @@ -142,12 +172,27 @@ def get_context_test_cases(): """ cases = [] sweep = get_mla_module_sweep_spec("trtllm") + head_counts = get_msa_head_counts("trtllm") # native // tp of the declared shards, not the shared head grid + # Declared budget for the Triton reference implementation + # (mla_module.yaml mla_module_trtllm.context.msa_triton_max_tokens): its + # masked-dense FP32 prefill is quadratic per sequence (24.5 s per forward + # at 64 heads / isl 32768 on rtxpro6000) and OOMs at b*isl=131072, so the + # SMs that only have it stop at 32K tokens; fmha_sm100 keeps the shared cap. + sm_version = get_sm_version() + max_tokens = sweep.context_max_tokens + triton_max_tokens = sweep.context_msa_triton_max_tokens + if triton_max_tokens and _msa_implementation(sm_version) == "triton": + max_tokens = min(max_tokens, triton_max_tokens) + dropped_triton_budget = 0 for compute_dtype, kv_dtype, gemm_type in _get_precision_combos("context"): - for num_heads in sweep.inner_sweep_head_counts: + for num_heads in head_counts: for b in sweep.context_batch_sizes: for s in sweep.context_sequence_lengths: if b * s > sweep.context_max_tokens: continue + if b * s > max_tokens: + dropped_triton_budget += len(sweep.context_prefix_lengths) + continue if ( sweep.context_large_sequence_min and s >= sweep.context_large_sequence_min @@ -156,6 +201,12 @@ def get_context_test_cases(): continue for prefix_len in sweep.context_prefix_lengths: cases.append([s, b, num_heads, kv_dtype, compute_dtype, gemm_type, prefix_len]) + if dropped_triton_budget: + print( + f"[trtllm-msa-module-cases] context: dropped {dropped_triton_budget} cases with batch*isl > " + f"{max_tokens} (SM{sm_version} runs the Triton reference implementation; " + "mla_module_trtllm.context.msa_triton_max_tokens)" + ) return cases @@ -167,8 +218,9 @@ def get_generation_test_cases(): """ cases = [] sweep = get_mla_module_sweep_spec("trtllm") + head_counts = get_msa_head_counts("trtllm") for compute_dtype, kv_dtype, gemm_type in _get_precision_combos("generation"): - for num_heads in sweep.inner_sweep_head_counts: + for num_heads in head_counts: for b in sweep.generation_batch_sizes: for s in sweep.generation_sequence_lengths: if b * s > sweep.generation_max_tokens: @@ -362,7 +414,7 @@ def _required_sparse_field(key: str) -> int: # elsewhere. kernel_source records which one actually ran. Field-presence # probe keeps rc19/rc20 (no such field, pydantic strict) working. if "implementation" in getattr(MiniMaxM3SparseAttentionConfig, "model_fields", {}): - sparse_kwargs["implementation"] = "msa" if get_sm_version() in (100, 103) else "triton" + sparse_kwargs["implementation"] = _msa_implementation(get_sm_version()) sparse_attention_config = MiniMaxM3SparseAttentionConfig(**sparse_kwargs) # NOTE: sparse_attention_config is intentionally NOT passed as a @@ -509,9 +561,14 @@ def _resolve_msa_metadata_cls(): instead of sniffing the version string. """ try: - from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_backend import ( - MiniMaxM3MsaSparseAttentionMetadata, - ) + try: + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( + MiniMaxM3MsaSparseAttentionMetadata, + ) + except ModuleNotFoundError: # < rc29 layout + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_backend import ( + MiniMaxM3MsaSparseAttentionMetadata, + ) except ImportError: # rc19/rc20 layout — Triton contract only return None return MiniMaxM3MsaSparseAttentionMetadata @@ -540,6 +597,28 @@ def _resolve_attention_backend_cls(model_config: ModelConfig): ) +# One MiniMax-M3 sparse attention layer per (model, heads, gemm type) per worker — +# same pattern and rationale as collect_mla_module._cached_attention_layer / +# collect_dsv4_attn._cached_dsv4_attention_module: the layer build does not depend +# on the point; the per-case KV cache manager, metadata, dry run and timing are +# untouched (b300 base run: msa_context_module 1.8 GPU-s per one-row case). +_LAYER_CACHE: dict = {} + + +def _cached_msa_attention_layer(model_path: str, num_heads: int, gemm_type: str, device: str): + key = (model_path, int(num_heads), gemm_type, device) + hit = _LAYER_CACHE.get(key) + if hit is not None: + return hit + if _LAYER_CACHE: + _LAYER_CACHE.clear() + gc.collect() + torch.cuda.empty_cache() + entry = create_msa_attention_layer(model_path=model_path, num_heads=num_heads, gemm_type=gemm_type, device=device) + _LAYER_CACHE[key] = entry + return entry + + def create_kv_cache_and_metadata( model_config: ModelConfig, batch_size: int, @@ -756,7 +835,15 @@ class (never the version string): "the MSA cache-write buffers (msa_backend.py:549-609@1.3.0rc23); " "the KV cache manager is not the M3 sparse manager" ) - if not is_context and attn_metadata.msa_decode_proxy_plan is None: + # rc23 exposed the graph-safe decode plan as ``msa_decode_proxy_plan``; rc29's + # MiniMaxM3MsaSparseAttentionMetadata no longer has that attribute (AttributeError + # on every generation case: 2,760/2,760, b200 job 469988228). Check the plan only + # where the contract still names it; on later series rely on _msa_fields_ready + # above. FIXME(re-verify): confirm rc29's decode-plan field name and restore an + # explicit check once known. + unset = object() + proxy_plan = getattr(attn_metadata, "msa_decode_proxy_plan", unset) + if not is_context and proxy_plan is None: raise RuntimeError( "MSA decode plans were not built for a pure generation batch " "(msa_backend.py:428-460@1.3.0rc23); sparse_metadata_params " @@ -817,11 +904,8 @@ def run_msa_module( f"compute={compute_dtype}, kv={kv_cache_dtype}, model={model_path}" ) - attn_module, model_config, original_architecture = create_msa_attention_layer( - model_path=model_path, - num_heads=num_heads, - gemm_type=gemm_type, - device=device, + attn_module, model_config, original_architecture = _cached_msa_attention_layer( + model_path, num_heads, gemm_type, device ) kv_cache_manager, attn_metadata, kernel_source = create_kv_cache_and_metadata( @@ -907,6 +991,10 @@ def kernel_func(): device=torch_device, kernel_func=kernel_func, num_warmups=warming_up, + # Graph-timed cases: the dry run already executed this forward, three + # eager passes suffice before capture. Eager-timed MSA prefill keeps + # the full warmup since it IS the measurement's warmup. + num_eager_warmups=3 if use_cuda_graph else warming_up, num_runs=test_ite, repeat_n=1, allow_graph_fail=False, @@ -1000,11 +1088,15 @@ def run_msa_module_worker( ) +# gc.collect() after every case cost 0.17 s of a ~1 s module case; collect periodically instead. +_CASE_GC = PeriodicGarbageCollector() + + def _cleanup(kv_cache_manager): if kv_cache_manager is not None: kv_cache_manager.shutdown() torch.cuda.empty_cache() - gc.collect() + _CASE_GC.maybe_collect() # ═══════════════════════════════════════════════════════════════════════ diff --git a/python/aisimulate/collector/trtllm/registry.py b/python/aisimulate/collector/trtllm/registry.py index 47589e145a..2405110ffc 100644 --- a/python/aisimulate/collector/trtllm/registry.py +++ b/python/aisimulate/collector/trtllm/registry.py @@ -24,6 +24,9 @@ get_func="get_computescale_test_cases", run_func="run_computescale", perf_filename=PerfFile.COMPUTESCALE, + # one measurement, two tables: the static quant row goes to + # scale_matrix_perf; finalize binds both to this producer + extra_perf_filenames=(PerfFile.SCALE_MATRIX,), ), OpEntry( op="mla_context", diff --git a/python/aisimulate/collector/version_resolver.py b/python/aisimulate/collector/version_resolver.py index 8dc208930b..eaa7abbec7 100644 --- a/python/aisimulate/collector/version_resolver.py +++ b/python/aisimulate/collector/version_resolver.py @@ -167,7 +167,8 @@ def build_collections( "get_func": entry.get_func, "run_func": entry.run_func, "perf_filename": entry.perf_filename, - "worker_perf_filename": entry.worker_perf_filename, + "extra_perf_filenames": tuple(entry.extra_perf_filenames), + "run_kwargs": tuple(entry.run_kwargs), "unverified": entry.unverified, "unverified_sms": entry.unverified_sms, } diff --git a/python/aisimulate/collector/vllm/collect_attn.py b/python/aisimulate/collector/vllm/collect_attn.py index 8bbde7e404..9e6e53eca4 100644 --- a/python/aisimulate/collector/vllm/collect_attn.py +++ b/python/aisimulate/collector/vllm/collect_attn.py @@ -1,13 +1,11 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 - -"""vLLM 0.24.0 dense-attention collector for CUDA backends.""" - -# B200 0.25.0 qualification (installed vLLM dd10e03f9), job 1967975: -# attention_context and attention_generation: 8/8 representative cases each. The native framework -# builders/selectors remain authoritative; no kernel fallback is introduced. -# The campaign manifest still selects one exact release per run. -__compat__ = "vllm>=0.24.0,<=0.25.0" +# +# Versions are upgraded IN PLACE (owner policy 2026-09-20): this file always +# targets the manifest pin; prior-version code is `git log -- ` +# away and prior-version DATA stays permanent under its version key. The +# 0.24->0.29 adaptations each carry a serving citation @0.29.0. +__compat__ = "vllm==0.30.0" import os @@ -24,6 +22,7 @@ create_kv_cache_and_block_mappings, create_standard_kv_cache_spec, create_vllm_config, + framework_kv_block_size, get_attention_backend, with_exit_stack, ) @@ -32,7 +31,10 @@ from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.utils.torch_utils import set_random_seed from vllm.v1.attention.backends.registry import AttentionBackendEnum -from vllm.v1.attention.backends.utils import set_kv_cache_layout +from vllm.v1.attention.backends.utils import ( + get_supported_kv_cache_layouts, + resolve_kv_cache_layout, +) from vllm.v1.attention.selector import AttentionSelectorConfig from vllm.version import __version__ as vllm_version @@ -54,6 +56,19 @@ def __init__(self, device: torch.device): # support MHA GQA MQA bfloat16 tensor and bfloat16/fp8 kv cache +def kv_cache_in_resolved_layout(kv_cache: torch.Tensor, layer_view_order: tuple[int, ...]) -> torch.Tensor: + """Return a tensor with the SAME logical [B, H, N, C] view as ``kv_cache`` + whose memory follows the resolved physical axis order ``layer_view_order`` + (vllm.v1.kv_cache_layout.KVCacheLayout.layer_view_order), i.e. exactly what + ``kv_cache.permute(*layer_view_order)`` in the backends expects to be + contiguous. Identity layout -> the input itself.""" + order = tuple(int(i) for i in layer_view_order) + if order == tuple(range(kv_cache.dim())): + return kv_cache.contiguous() + inverse = [order.index(i) for i in range(len(order))] + return kv_cache.permute(*order).contiguous().permute(*inverse) + + def _dense_kernel_source(backend_name_str, impl, attn_metadata): """Ground-truth kernel_source for the dense attention row.""" if backend_name_str == "FLASH_ATTN": @@ -73,7 +88,17 @@ def _dense_kernel_source(backend_name_str, impl, attn_metadata): phase_metadata = attn_metadata.decode if phase_metadata is None: raise RuntimeError("vLLM FlashInfer metadata has neither a prefill nor a decode portion") - return f"vllm_flashinfer_{type(phase_metadata).__name__}".lower() + label = f"vllm_flashinfer_{type(phase_metadata).__name__}".lower() + # 0.30 folds XQA and trtllm-gen decode into ONE metadata class, + # FlashInferTrtllmAPIDecode, and names the kernel in its `kernel` field + # (flashinfer.py:568-572 FlashInferDecodeKernel, :608-611 @v0.30.0; + # decode_with_xqa / decode_with_trtllm_gen at :1997-1998). The class + # name alone would label SM120's XQA row and SM100's trtllm-gen row + # identically — kernel_source must name the kernel that ran. + kernel = getattr(phase_metadata, "kernel", None) + if kernel is not None: + label += "_" + str(getattr(kernel, "value", kernel)).lower().replace("-", "_") + return label return f"vllm_{backend_name_str}".lower() @@ -96,7 +121,10 @@ def run_attention_torch( dtype = torch.bfloat16 model = os.path.join(os.path.dirname(__file__), "fake_hf_model") - block_size = 64 + # Backend selection runs at the framework default page (serving does the same: the + # model is built at CacheConfig.DEFAULT_BLOCK_SIZE, the page is re-derived from the + # selected backend afterwards — utils.framework_kv_block_size). + block_size = None # unconstrained selection (no --block-size); the page is asked from the backend below if is_context_phase: batch_spec = BatchSpec( @@ -132,14 +160,35 @@ def run_attention_torch( head_size=head_dim, dtype=dtype, kv_cache_dtype="fp8" if use_fp8_kv_cache else "auto", - block_size=block_size, + block_size=None, # serving passes None unless the user set --block-size (selector.py:131-135) use_mla=False, has_sink=False, use_sparse=False, use_mm_prefix=False, ) + # Backend selection for head_dim > 256 mirrors the framework's own + # heterogeneous-head-dim policy. head_dim > 256 is reachable only from + # Gemma-4 (global layers, global_head_dim=512), whose model-level config + # hook (Gemma4Config.verify_and_update_config, + # model_executor/models/config.py:209-256) forces a UNIFORM backend across + # its mixed 256/512 heads to avoid the FA3/FA4 mixed penalty. That policy, + # plus the per-group validity of head-512, resolves as follows on SM90 + # (live-engine kernel probe, 2026-09-22, gemma-4-26B-A4B + kv fp8/auto): + # bf16-KV -> FLASH_ATTN (FA4 cute flash_fwd_sm90) — the bare selector + # already returns this, so no override is needed. + # fp8-KV -> TRITON_ATTN (kernel_unified_attention). No FA/FlashInfer + # path compiles fp8+head512 on SM90, so the uniform-backend + # policy lands on Triton. The bare platform selector, lacking + # the model-level heterogeneous-head context, instead picks + # FlashInfer (whose SM90 fp8+head512 JIT cannot compile) — the + # attention fp8 collection failures. Force the serving backend. + # This is not a per-model special case: it reproduces the framework's own + # config-hook branch keyed on head_dim, with the metadata-parity citation. + selected_backend = None + if head_dim > 256 and use_fp8_kv_cache: + selected_backend = AttentionBackendEnum.TRITON_ATTN backend_path = current_platform.get_attn_backend_cls( - None, + selected_backend, attn_selector_config, num_heads=num_heads, ) @@ -147,6 +196,10 @@ def run_attention_torch( backend_name_str = backend_cls.get_name() backend_name = AttentionBackendEnum[backend_name_str] + # the page the framework would run THIS backend with (asked, not tabled) + block_size = framework_kv_block_size(backend_cls) + vllm_config.cache_config.block_size = block_size + kv_cache_spec = create_standard_kv_cache_spec(vllm_config, use_fp8_kv_cache) # Ensure the KV cache has enough blocks for all sequences. @@ -195,15 +248,33 @@ def run_attention_torch( randomize_blocks=True, ) - # The helper populates [2, num_blocks, ...]; vLLM 0.24 uses logical - # [num_blocks, 2, ...] and may impose a backend-specific physical stride. - kv_cache = kv_cache.transpose(0, 1).contiguous() - set_kv_cache_layout(backend_cls.get_required_kv_cache_layout()) - exit_stack.callback(set_kv_cache_layout, None) - stride_order = backend_cls.get_kv_cache_stride_order() - if stride_order != tuple(range(kv_cache.ndim)): - inverse_order = [stride_order.index(i) for i in range(kv_cache.ndim)] - kv_cache = kv_cache.permute(*stride_order).contiguous().permute(*inverse_order) + # serving-same layout resolution for the single backend under test + # (utils.py:240-307@v0.29.0); reset cache_config after the case so the + # next backend re-resolves instead of inheriting ("existing wins outright") + _supported = get_supported_kv_cache_layouts([backend_cls]) + resolve_kv_cache_layout(vllm_config, [[m.name for m in _supported]], [kv_cache_spec]) + exit_stack.callback(lambda: setattr(vllm_config.cache_config, "kv_cache_layout", None)) + # 0.29 impls consume ONE merged tensor [num_blocks, H, block, 2*D] with + # K/V split on the last dim (flash_attn.py:1013-1014 "(B,H,N,2*D) -> + # ((B,N,H,D),(B,N,H,D))"); the 0.24-era per-backend stride permutation + # (get_kv_cache_stride_order) no longer exists. Convert the populated + # helper layout [2, B, N, H, D] once, value-equivalent. + _k, _v = kv_cache[0], kv_cache[1] # [B, N, H, D] + kv_cache = torch.cat((_k.permute(0, 2, 1, 3), _v.permute(0, 2, 1, 3)), dim=-1).contiguous() # [B, H, N, 2*D] + # Serving allocates the per-layer cache in the RESOLVED physical layout and + # backends read it through kv_cache.permute(*layout.layer_view_order) + # (v1/kv_cache_layout.py KVCacheLayout.layer_view_order; flashinfer.py:2091 + # "kv_cache_permute = kv_cache.permute(*stride_order) # HND and contiguous" + # @v0.30.0). A contiguous [B, H, N, 2*D] tensor is only right for the LBHNC + # family (HND) that FLASH_ATTN and the SM100 trtllm-gen FlashInfer path + # resolve (flashinfer.py:543-549); FlashInfer elsewhere resolves the NHD + # default, whose XQA decode asserts "KV cache inner dims (block_size, + # head_size) must be contiguous" on this HND memory (flashinfer.py:2448) — + # every fp8-KV GQA generation case on SM120, 2026-10-01. Lay the memory out + # per the resolved layout and hand the backend the same logical view. + kv_cache = kv_cache_in_resolved_layout( + kv_cache, vllm_config.cache_config.get_resolved_kv_cache_layout().layer_view_order + ) builder_cls, impl_cls = get_attention_backend(backend_name) layer_names = ["placeholder"] diff --git a/python/aisimulate/collector/vllm/collect_attn_encoder.py b/python/aisimulate/collector/vllm/collect_attn_encoder.py index 7537a0612f..cb96760a62 100644 --- a/python/aisimulate/collector/vllm/collect_attn_encoder.py +++ b/python/aisimulate/collector/vllm/collect_attn_encoder.py @@ -14,7 +14,7 @@ # encoder_attention: 8/8 representative cases. The native framework # builders/selectors remain authoritative; no kernel fallback is introduced. # The campaign manifest still selects one exact release per run. -__compat__ = "vllm>=0.24.0,<=0.25.0" +__compat__ = "vllm==0.30.0" import torch from collector.case_generator import get_attention_encoder_head_configs, get_attention_encoder_shape_sweeps diff --git a/python/aisimulate/collector/vllm/collect_computescale.py b/python/aisimulate/collector/vllm/collect_computescale.py index 905b79232c..48c07d5a23 100644 --- a/python/aisimulate/collector/vllm/collect_computescale.py +++ b/python/aisimulate/collector/vllm/collect_computescale.py @@ -1,13 +1,20 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Measure vLLM FP8 activation quantization overhead for static-FP8 GEMM.""" +"""Measure vLLM FP8 activation quantization overhead for static-FP8 GEMM. + +One measurement, two tables — the registry entry declares the second one as +``extra_perf_filenames`` and the executor hands both names in, so finalize +binds both to this producer's checkpoint: + perf_filename -> computescale_perf.txt (dynamic per-token quant minus static quant) + extra_perf_filenames[0] -> scale_matrix_perf.txt (the static per-tensor quant alone) +""" # B200 0.25.0 qualification (installed vLLM dd10e03f9), job 1968047: # compute_scale: 8/8 representative cases. The native framework # builders/selectors remain authoritative; no kernel fallback is introduced. # The campaign manifest still selects one exact release per run. -__compat__ = "vllm>=0.24.0,<=0.25.0" +__compat__ = "vllm==0.30.0" import torch from collector.case_generator import get_compute_scale_case_specs @@ -17,6 +24,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape from vllm.version import __version__ as vllm_version +_OUTSIDE_LOOP_COUNT = 5 + def get_computescale_test_cases(): if get_sm_version() <= 86: @@ -24,37 +33,46 @@ def get_computescale_test_cases(): return [[case.m, case.k] for case in get_compute_scale_case_specs()] -def run_computescale(m, k, *, perf_filename, device="cuda:0"): +def _setup(m, k, device): setup_distributed(device) device = torch.device(device) torch.cuda.set_device(device) torch.set_default_device(device) - x = torch.randn((m, k), dtype=torch.bfloat16, device=device) - static_scale = torch.tensor([1.0], dtype=torch.float32, device=device) - outside_loop_count = 5 + return device, x + - def dynamic_kernel_func(): - for _ in range(outside_loop_count): +def _bench_dynamic(device, x): + def kernel_func(): + for _ in range(_OUTSIDE_LOOP_COUNT): ops.scaled_fp8_quant(x, scale=None, use_per_token_if_dynamic=True) - with benchmark_with_power(device=device, kernel_func=dynamic_kernel_func, repeat_n=1) as dynamic_results: + with benchmark_with_power(device=device, kernel_func=kernel_func, repeat_n=1) as results: pass + return results["latency_ms"] / _OUTSIDE_LOOP_COUNT, results["power_stats"] - dynamic_latency = dynamic_results["latency_ms"] / outside_loop_count - def static_kernel_func(): - for _ in range(outside_loop_count): +def _bench_static(device, x): + static_scale = torch.tensor([1.0], dtype=torch.float32, device=device) + + def kernel_func(): + for _ in range(_OUTSIDE_LOOP_COUNT): ops.scaled_fp8_quant( x, scale=static_scale, group_shape=(GroupShape.PER_TENSOR.row, GroupShape.PER_TENSOR.col), ) - with benchmark_with_power(device=device, kernel_func=static_kernel_func, repeat_n=1) as static_results: + with benchmark_with_power(device=device, kernel_func=kernel_func, repeat_n=1) as results: pass + return results["latency_ms"] / _OUTSIDE_LOOP_COUNT, results["power_stats"] - static_latency = static_results["latency_ms"] / outside_loop_count + +def run_computescale(m, k, *, perf_filename, extra_perf_filenames, device="cuda:0"): + (scale_matrix_filename,) = extra_perf_filenames + device, x = _setup(m, k, device) + dynamic_latency, dynamic_power = _bench_dynamic(device, x) + static_latency, static_power = _bench_static(device, x) compute_scale_latency = max(0.0, dynamic_latency - static_latency) log_perf( @@ -65,8 +83,9 @@ def static_kernel_func(): op_name="compute_scale", kernel_source="dynamic_per_token_scaled_fp8_quant_minus_static_scaled_fp8_quant", perf_filename=perf_filename, - power_stats=dynamic_results["power_stats"], + power_stats=dynamic_power, ) + log_perf( item_list=[{"m": m, "k": k, "quant_dtype": "fp8", "latency": static_latency}], framework="VLLM", @@ -74,6 +93,6 @@ def static_kernel_func(): device_name=torch.cuda.get_device_name(device), op_name="scale_matrix", kernel_source="static_scaled_fp8_quant", - perf_filename="scale_matrix_perf.txt", - power_stats=static_results["power_stats"], + perf_filename=scale_matrix_filename, + power_stats=static_power, ) diff --git a/python/aisimulate/collector/vllm/collect_dsv4_attn.py b/python/aisimulate/collector/vllm/collect_dsv4_attn.py index 8146fcda48..0f007b22d5 100644 --- a/python/aisimulate/collector/vllm/collect_dsv4_attn.py +++ b/python/aisimulate/collector/vllm/collect_dsv4_attn.py @@ -42,7 +42,6 @@ _DSV4_SPARSE_ISL_LIST, _DSV4_SPARSE_MAX_FULL_S, _DSV4_SPARSE_PAST_KV_LIST, - _DSV4_SPARSE_TP_LIST_ATTN, _DSV4_SPARSE_TP_LIST_INDEXER, DSV4_ATTN_KINDS, DSV4_SPARSE_KERNELS, @@ -73,21 +72,22 @@ # CSA/HCA context/generation: 32/32; auxiliary sparse kernels: 4/4. The native framework # builders/selectors remain authoritative; no kernel fallback is introduced. # The campaign manifest still selects one exact release per run. -__compat__ = "vllm>=0.24.0,<=0.25.0" +# Versions are upgraded IN PLACE (owner policy 2026-09-20): this file always +# targets the manifest pin; prior-version code is `git log -- ` +# away and prior-version DATA stays permanent under its version key. The +# 0.24->0.29 adaptations each carry a serving citation @0.29.0. +__compat__ = "vllm==0.30.0" DEFAULT_MODEL = _DSV4_DEFAULT_MODELS[0] ARCHITECTURE = "DeepseekV4ForCausalLM" ATTN_KIND_TO_COMPRESS_RATIO = {"csa": 4, "hca": 128} -SPARSE_KERNEL_TO_ATTN_KIND = {"paged_mqa_logits": "csa", "hca_attn": "hca"} -SPARSE_KERNEL_TO_OP_NAME = { - "paged_mqa_logits": "dsv4_paged_mqa_logits_module", - "hca_attn": "dsv4_hca_attn_module", -} -SPARSE_KERNEL_TO_PERF_FILE = { - "paged_mqa_logits": PerfFile.DSV4_PAGED_MQA_LOGITS_MODULE, - "hca_attn": PerfFile.DSV4_HCA_ATTN_MODULE, -} +# paged_mqa_logits (the CSA indexer scoring GEMM) is the only standalone sparse +# sub-kernel; the hca_attn FMLA sidecar was removed 2026-10-15 (retired SDK +# placeholder, see case_generator.DSV4_SPARSE_KERNELS). +SPARSE_KERNEL_TO_ATTN_KIND = {"paged_mqa_logits": "csa"} +SPARSE_KERNEL_TO_OP_NAME = {"paged_mqa_logits": "dsv4_paged_mqa_logits_module"} +SPARSE_KERNEL_TO_PERF_FILE = {"paged_mqa_logits": PerfFile.DSV4_PAGED_MQA_LOGITS_MODULE} MODEL_CONFIGS_DIR = Path(__file__).resolve().parents[2] / "src" / "aisimulate_core" / "model_configs" SUPPORTED_GEMM_TYPES = {"fp8_block"} DEFAULT_MAX_SEQ_LEN = 65536 @@ -253,7 +253,7 @@ def _create_dsv4_attention_module( capability = current_platform.get_device_capability() if capability is not None and capability.to_int() == 89: raise RuntimeError( - "vLLM 0.24.0 DeepSeek-V4 attention is unsupported on SM89: " + "vLLM DeepSeek-V4 attention is unsupported on SM89 (capabilities.yaml op_min_sm floors it at 90): " "the production selector falls back to FlashMLA, whose DSV4 " "backend supports SM90 and SM10x only." ) @@ -361,23 +361,55 @@ def _remap_common_metadata(common, *, block_size: int, device: str): def _allocate_attention_kv_cache(backend, spec, num_blocks: int, cache_dtype: str, *, device: str) -> torch.Tensor: - shape_block_size = spec.storage_block_size if spec.storage_block_size != spec.block_size else spec.block_size - shape = backend.get_kv_cache_shape( - num_blocks, - shape_block_size, - spec.num_kv_heads, - spec.head_size, - cache_dtype, - ) + # 0.29: the shape contract lives in the generic spec pipeline. Serving + # allocates flat bytes and views each layer as the 4D logical [B,H,N,C] + # page (compute_layer_kv_cache_shape_bytes, v1/kv_cache_interface.py:253-269; + # per-layer element view = bytes.view(spec.dtype), create_kv_cache_views). + # The dsv4 layout facts (584B fp8_ds_mla state content, 576B alignment, + # tokens_per_state=compress_ratio) are spec fields at 0.29 + # (models/deepseek_v4/attention.py:741-753), so the helper output IS the + # serving shape. Page padding: mirror create_kv_cache_views' block stride + # (page_size_bytes per block) via empty_strided when the spec pads pages. + from vllm.v1.kv_cache_interface import compute_layer_kv_cache_shape_bytes + + shape_bytes = compute_layer_kv_cache_shape_bytes(spec, num_blocks) + dtype_size = torch.empty((), dtype=spec.dtype).element_size() + shape = (*shape_bytes[:-1], shape_bytes[-1] // dtype_size) if spec.page_size_padded is None: return torch.zeros(shape, dtype=spec.dtype, device=device) - dtype_size = torch.empty((), dtype=spec.dtype).element_size() strides = list(torch.empty(shape).stride()) strides[0] = spec.page_size_bytes // dtype_size return torch.empty_strided(shape, tuple(strides), dtype=spec.dtype, device=device).zero_() +def _make_builder(backend, spec, prefix, vllm_config, common, *, device: str): + """Construct the layer's metadata builder exactly as serving does at 0.29. + + Builders may demand kw-only block_table_width (DeepseekV32IndexerMetadataBuilder, + mla/indexer.py:519,531); serving computes it via get_block_table_width and + ALLOCATES the block table at that width (v1/worker/utils.py:276-283 + + v1/worker/block_table.py:29-49), so builders assume + table.shape[1] == width — the family-100 decode path copies the table into + a width-sized expanded buffer (indexer.py:609,745) and crashes on narrower + tables (found on B300 2026-09-20 in the mla 029 lane; same contract here). + Zero-pad our ceil(seq/block)-column table to the width. + """ + builder_cls = backend.get_builder_cls() + kwargs = {} + if getattr(builder_cls, "requires_block_table_width", False): + from vllm.v1.worker.block_table import get_block_table_width + + max_num_blocks = spec.max_num_blocks_per_req(vllm_config, vllm_config.model_config.max_model_len) + width = get_block_table_width(max_num_blocks, spec.block_size) + kwargs["block_table_width"] = width + _bt = common.block_table_tensor + if _bt.shape[1] < width: + _pad = torch.zeros((_bt.shape[0], width - _bt.shape[1]), dtype=_bt.dtype, device=_bt.device) + common.block_table_tensor = torch.cat([_bt, _pad], dim=1) + return builder_cls(spec, [prefix], vllm_config, torch.device(device), **kwargs) + + def _build_metadata_and_bind_caches(attn_module: DeepseekV4Attention, vllm_config, common, *, device: str): cache_blocks = _cache_blocks(int(common.num_reqs), int(common.max_seq_len)) metadata = {} @@ -392,19 +424,21 @@ def _build_metadata_and_bind_caches(attn_module: DeepseekV4Attention, vllm_confi if spec is None: continue backend = registered_layer.get_attn_backend() - metadata[layer.prefix] = backend.get_builder_cls()( - spec, - [layer.prefix], - vllm_config, - torch.device(device), - ).build(0, common) + metadata[layer.prefix] = _make_builder(backend, spec, layer.prefix, vllm_config, common, device=device).build( + 0, common + ) cache_dtype = getattr(spec, "cache_dtype_str", None) or "auto" - registered_layer.kv_cache = _allocate_attention_kv_cache( - backend, - spec, - cache_blocks, - cache_dtype, - device=device, + # serving binds through each layer's bind_kv_cache + # (v1/worker/utils.py:622@0.29.0) — layer overrides may reshape the + # raw view (e.g. squeeze the H=1 dim), so never assign .kv_cache raw + registered_layer.bind_kv_cache( + _allocate_attention_kv_cache( + backend, + spec, + cache_blocks, + cache_dtype, + device=device, + ) ) compressors = [attn_module.compressor] @@ -415,23 +449,22 @@ def _build_metadata_and_bind_caches(attn_module: DeepseekV4Attention, vllm_confi spec = state_cache.get_kv_cache_spec(vllm_config) backend = state_cache.get_attn_backend() compressor_common = _remap_common_metadata(common, block_size=spec.block_size, device=device) - metadata[state_cache.prefix] = backend.get_builder_cls()( - spec, - [state_cache.prefix], - vllm_config, - torch.device(device), + metadata[state_cache.prefix] = _make_builder( + backend, spec, state_cache.prefix, vllm_config, compressor_common, device=device ).build(0, compressor_common) state_cache_blocks = _cache_blocks_for_block_size( int(common.num_reqs), int(common.max_seq_len), spec.block_size, ) - state_cache.kv_cache = _allocate_attention_kv_cache( - backend, - spec, - state_cache_blocks, - getattr(spec, "cache_dtype_str", None) or "auto", - device=device, + state_cache.bind_kv_cache( + _allocate_attention_kv_cache( + backend, + spec, + state_cache_blocks, + getattr(spec, "cache_dtype_str", None) or "auto", + device=device, + ) ) return metadata @@ -674,46 +707,6 @@ def kernel_func(): return result -def _bench_mla_sparse_op( - *, - attn_module: DeepseekV4Attention, - vllm_config, - metadata, - positions: torch.Tensor, - num_tokens: int, - device: str, - warming_up: int, - test_ite: int, -) -> dict: - sparse_attn = attn_module - q = torch.full( - (num_tokens, sparse_attn.padded_heads, sparse_attn.head_dim), - 0.01, - dtype=torch.bfloat16, - device=device, - ) - kv = torch.zeros((num_tokens, sparse_attn.head_dim), dtype=torch.bfloat16, device=device) - output = torch.empty_like(q) - - with set_current_vllm_config(vllm_config), set_forward_context(metadata, vllm_config), torch.inference_mode(): - sparse_attn.forward_mqa(q, kv, positions, output) - torch.cuda.synchronize() - - def kernel_func(): - sparse_attn.forward_mqa(q, kv, positions, output) - - with benchmark_with_power( - device=torch.device(device), - kernel_func=kernel_func, - num_warmups=warming_up, - num_runs=test_ite, - repeat_n=1, - use_cuda_graph=True, - ) as result: - pass - return result - - def _bench_sparse_kernel_shape( *, model_path: str, @@ -734,64 +727,21 @@ def _bench_sparse_kernel_shape( raise ValueError(f"invalid sparse sequence length: isl={isl}, past_kv={past_kv}") attn_kind = SPARSE_KERNEL_TO_ATTN_KIND[kernel] - is_context = isl > 1 - query_len = isl - num_tokens = batch_size * query_len - - if kernel == "paged_mqa_logits": - hf_config = SimpleNamespace(**_read_model_config(model_path)) - architecture = hf_config.architectures[0] if hf_config.architectures else ARCHITECTURE - native_num_heads = int(hf_config.num_attention_heads) - local_num_heads = native_num_heads // tp_size - backend_name = "vllm.utils.deep_gemm.fp8_fp4_paged_mqa_logits" - result = _bench_paged_mqa_logits_kernel( - hf_config=hf_config, - num_query_rows=num_tokens, - past_kv=past_kv, - device=device, - warming_up=warming_up, - test_ite=test_ite, - ) - else: - with ( - _tp_simulation(tp_size), - _create_dsv4_attention_module( - model_path=model_path, - attn_kind=attn_kind, - batch_size=batch_size, - seq_len=full_seq_len, - tp_size=tp_size, - is_context=is_context, - device=device, - query_len=query_len, - ) as (attn_module, vllm_config), - ): - common = _make_common_metadata( - batch_size=batch_size, - seq_len=full_seq_len, - is_context=is_context, - query_len=query_len, - device=device, - ) - metadata = _build_metadata_and_bind_caches(attn_module, vllm_config, common, device=device) - positions = common.positions - hf_config = vllm_config.model_config.hf_config - architecture = hf_config.architectures[0] if hf_config.architectures else ARCHITECTURE - native_num_heads = int(hf_config.num_attention_heads) - local_num_heads = int(attn_module.n_local_heads) - concrete_layer = vllm_config.compilation_config.static_forward_context[attn_module.prefix] - backend_name = concrete_layer.get_attn_backend().get_name() - result = _bench_mla_sparse_op( - attn_module=attn_module, - vllm_config=vllm_config, - metadata=metadata, - positions=positions, - num_tokens=num_tokens, - device=device, - warming_up=warming_up, - test_ite=test_ite, - ) - del attn_module, vllm_config, metadata, common, positions + num_tokens = batch_size * isl + + hf_config = SimpleNamespace(**_read_model_config(model_path)) + architecture = hf_config.architectures[0] if hf_config.architectures else ARCHITECTURE + native_num_heads = int(hf_config.num_attention_heads) + local_num_heads = native_num_heads // tp_size + backend_name = "vllm.utils.deep_gemm.fp8_fp4_paged_mqa_logits" + result = _bench_paged_mqa_logits_kernel( + hf_config=hf_config, + num_query_rows=num_tokens, + past_kv=past_kv, + device=device, + warming_up=warming_up, + test_ite=test_ite, + ) latency = float(result["latency_ms"]) @@ -800,7 +750,7 @@ def _bench_sparse_kernel_shape( { "model": model_path, "architecture": architecture, - "mla_dtype": "fp8_e4m3" if kernel == "paged_mqa_logits" else "bfloat16", + "mla_dtype": "fp8_e4m3", "kv_cache_dtype": "fp8", "gemm_type": "fp8_block", # Sparse-op consumers key ``num_heads`` by the native model @@ -913,7 +863,7 @@ def _build_vllm_dsv4_sparse_test_cases(kernel: str) -> list[list]: if not model_paths: return [] - tp_list = _DSV4_SPARSE_TP_LIST_ATTN if kernel == "hca_attn" else _DSV4_SPARSE_TP_LIST_INDEXER + tp_list = _DSV4_SPARSE_TP_LIST_INDEXER if "--smoke" in sys.argv: smoke_shapes = [ (1, 1, 8192), @@ -1058,10 +1008,6 @@ def get_dsv4_paged_mqa_logits_test_cases(): return _build_vllm_dsv4_sparse_test_cases("paged_mqa_logits") -def get_dsv4_hca_attn_test_cases(): - return _build_vllm_dsv4_sparse_test_cases("hca_attn") - - def main() -> None: parser = argparse.ArgumentParser(description="Collect vLLM DeepSeek-V4 attention module latency.") parser.add_argument("--model-path", default=DEFAULT_MODEL) diff --git a/python/aisimulate/collector/vllm/collect_gdn.py b/python/aisimulate/collector/vllm/collect_gdn.py index 6a4e468a9d..4bb9bfb845 100644 --- a/python/aisimulate/collector/vllm/collect_gdn.py +++ b/python/aisimulate/collector/vllm/collect_gdn.py @@ -60,7 +60,7 @@ # 1968047). Existing framework construction/dispatch is unchanged. Invalid # model/quant shapes continue to raise and remain recorded failures; this is # API compatibility, not a claim that every Cartesian shape is supported. -__compat__ = "vllm>=0.24.0,<=0.27.1,!=0.25.1,!=0.26.0,!=0.27.0" +__compat__ = "vllm==0.30.0" import gc import os @@ -369,6 +369,10 @@ def run_conv1d(_conv_input=conv_input, _conv_state=conv_state): # present and source-unchanged, so the claim is not dangling, but # the bug reproduction itself is unconfirmed at this version. # Re-verify on the next vLLM bump or the first 0.27.1 GPU run. + # RE-VERIFIED on hardware at 0.30.0, SM89 (L40, 2026-10-04, clean GPU, container per op): + # a 40-case --shuffle sample of the gdn op failed 21 tasks, 17 of them CUDA faults + # (9 illegal memory access + 8 SIGABRT) in this same chunked prefill at the large-token + # sub-points, 3 the grid-y limit below, 1 capacity OOM — the family is unchanged from 0.24.0. def run_gdn_scan(_q=q, _k=k, _v=v, _g=g, _beta=beta, _state=gdn_state): chunk_gdn( q=_q, diff --git a/python/aisimulate/collector/vllm/collect_gemm.py b/python/aisimulate/collector/vllm/collect_gemm.py index 0d5a94908e..35d47f6d26 100644 --- a/python/aisimulate/collector/vllm/collect_gemm.py +++ b/python/aisimulate/collector/vllm/collect_gemm.py @@ -58,9 +58,11 @@ # 322-363); NVFP4 still uses CT's factory and records the selected kernel # (schemes/compressed_tensors_w4a4_nvfp4.py:29-31,95-141). This adds the # exact 0.25.0 release to the existing lane, not 0.25.1/0.26.0/0.27.0. -# 0.30.0 source audit (2026-09-30: tag v0.30.0 == ced6857afa0e vs v0.27.1), -# the baseline of the 0.31.0 audit below. 0.28.0, 0.29.0 and 0.30.0 itself -# are not admitted: no consumer reads their tables. +# 0.30.0 source audit (2026-09-30: tag v0.30.0 == ced6857afa0e vs v0.27.1): +# the 2026-10 campaign tables and the baseline of the 0.31.0 audit below. +# The gemm family runs on 0.31.0 (framework_manifest.yaml family pin, +# 2026-10-30), so this file targets exactly 0.31.0. +# One release per file (owner policy 2026-09-20); earlier code is `git log -- ` away. # bf16: UnquantizedLinearMethod now binds its GEMM at __init__ from # kernel_config.linear_backend (layers/linear.py:169-174, apply :234-238); # dispatch_unquantized_gemm (layers/utils.py:584-616) returns @@ -128,8 +130,8 @@ # 0.7.0.post1 (was 0.6.18.post1). # Helpers with identical bodies: scaled_fp4_quant (_custom_ops.py:1513-1588), # per_block_cast_to_fp8 (utils/deep_gemm.py:983-1002). All imports below -# resolve at 0.31.0. 0.28.0/0.29.0 are still not audited. -__compat__ = "vllm>=0.24.0,<=0.31.0,!=0.25.1,!=0.26.0,!=0.27.0,!=0.28.0,!=0.29.0,!=0.30.0" +# resolve at 0.31.0. +__compat__ = "vllm==0.31.0" from types import SimpleNamespace @@ -204,6 +206,21 @@ def get_gemm_test_cases(): @with_exit_stack def run_gemm(exit_stack, gemm_type, m, n, k, *, perf_filename, device="cuda:0"): + if gemm_type == "fp8_block" and (n % 16 or k % 16) and get_sm_version() >= 90: + # FIXME(kernel-limit): vLLM routes block-fp8 GEMMs whose weight dims are not multiples + # of 16 away from cutlass (_custom_ops.cutlass_scaled_mm: `cutlass_compatible_b = + # b.shape[0] % 16 == 0 and b.shape[1] % 16 == 0` @0.31.0:807) into triton_scaled_mm, + # whose scale-shape assert cannot take 128x128 block scales (triton_scaled_mm.py:201). + # Deterministic per shape: gb200 1,110 cases, every one N in {1,4,8,12,24} (pipeline + # 72343556; b200/b300 ~130 of the same shard). The ed9109ec attempt to classify by a + # pre-flight forward never ran: create_gemm()'s own dry run asserts first. Refuse before + # the layer is built; serving linears with these N take the same path. SM89 never + # reaches it: the cutlass block tier needs SM90 and Marlin is selected first + # (_POSSIBLE_FP8_BLOCK_KERNELS, model_executor/kernels/linear/__init__.py:458-467 @0.31.0). + raise RuntimeError( + f"FIXME(kernel-limit): vLLM fp8_block GEMM needs n and k multiples of 16 for cutlass_scaled_mm " + f"(n={n}, k={k}); the triton_scaled_mm fallback asserts on 128x128 block scales" + ) setup_distributed(device) if envs.VLLM_BATCH_INVARIANT: diff --git a/python/aisimulate/collector/vllm/collect_kda.py b/python/aisimulate/collector/vllm/collect_kda.py index e5dcc0bbf9..d3c58dc506 100644 --- a/python/aisimulate/collector/vllm/collect_kda.py +++ b/python/aisimulate/collector/vllm/collect_kda.py @@ -48,21 +48,17 @@ kda_perf.txt — same column layout as the sglang kda collector. """ -# Serve-parity `metadata=` argument validated on the 0.27.0 release image -# (Kimi-K3 landed upstream there). The manifest kda family pin stays on the -# kimi-k3 preview build until the SDK consumer routes 0.27.0's split decode -# (see collector/framework_manifest.yaml), so this module exactly pins the -# preview version; it runs on either image — the DS-layout probe -# (is_fused_kda_decode_supported) yields fused_kda_decode rows on the -# preview and the packed conv-update + recurrence pair on 0.27.0. -# The file-level range spans the two audited routes (PEP 440 clauses are -# conjunctive, so it cannot name the two releases alone). Each KDA -# architecture is additionally gated at runtime to its own audited release by -# _KDA_ARCHITECTURE_COMPAT: Kimi-K3 only on the 0.1.dev19262 preview, GLM-5.3- -# Flash only on stock 0.31.0 (no overlay); every release in between -# (including 0.30.0, where the GLM layer lived under glm5next/nvidia/) raises -# KdaRuntimeNotAuditedError. -__compat__ = "vllm>=0.1.dev19262,<=0.31.0" +# The kda family runs on 0.31.0 (framework_manifest.yaml family pin, +# 2026-10-30) for both architectures. History: Kimi-K3 ran on 0.30.0 after +# the kimi-k3 preview pin was retired 2026-09-28: 0.30's DS-layout probe +# (is_fused_kda_decode_supported, now with recurrent_state_dtype) selects +# fused_kda_decode for the non-spec decode, and the packed conv-update + +# recurrence pair is measured alongside as the spec-decode serving path. +# GLM-5.3-Flash runs on stock 0.31.0 (no overlay). Each KDA architecture is +# additionally gated at runtime by _KDA_ARCHITECTURE_COMPAT, raising +# KdaRuntimeNotAuditedError otherwise. +# One release per file (owner policy 2026-09-20); earlier code is `git log -- ` away. +__compat__ = "vllm==0.31.0" import gc import os @@ -98,6 +94,20 @@ # (config.json linear_attn_config.gate_lower_bound). KDA_LOWER_BOUND = -5.0 + +def _decode_paths() -> set[str]: + """Which decode paths the generation benchmark measures. Default: both the + fused CUDA decode kernel (serving's non-spec path) and the packed + conv-update + Triton recurrence pair (serving's path under speculative + decoding). The opharness kda decode gate captures serving WITHOUT spec + decoding, so it sets AIS_KDA_DECODE_PATHS=fused to compare like with like + (recorded in the capture's env). Collection runs never set it.""" + paths = {p.strip() for p in os.environ.get("AIS_KDA_DECODE_PATHS", "fused,packed").split(",") if p.strip()} + if not paths or not paths <= {"fused", "packed"}: + raise ValueError(f"AIS_KDA_DECODE_PATHS must name fused and/or packed, got {sorted(paths)}") + return paths + + NUM_WARMUPS = 3 NUM_RUNS = 10 @@ -151,10 +161,16 @@ def _log(common, latency_ms, kernel_source, perf_filename, vllm_version, device, def _resolve_prefill_kernel(dtype: torch.dtype): """Mirror serving's resolve_kda_prefill_backend on this device: FlashKDA when supported AND importable, else the Triton chunk kernel - (vllm/models/kimi_k3/nvidia/kda.py resolve_kda_prefill_backend).""" + (vllm/models/kimi_k3/nvidia/kda.py:455-495@v0.30.0 resolve_kda_prefill_backend; + backend "auto" never selects FlashInfer, so on SM90 the choice is FlashKDA + or Triton). 0.30 added the recurrent-state dtype to the probe: serving + passes get_state_dtype()[1] (kda.py:665), which + MambaStateDtypeCalculator.kda_state_dtype resolves to float32 for the + default mamba_ssm_cache_dtype "auto" — the dtype this collector's + init_state uses below.""" from vllm.models.kimi_k3.nvidia.kda import is_flashkda_supported - if is_flashkda_supported(128, dtype, KDA_LOWER_BOUND): + if is_flashkda_supported(128, dtype, torch.float32, KDA_LOWER_BOUND): try: import vllm._flashkda_C # noqa: F401 from vllm.models.kimi_k3.nvidia.kda import _flashkda_prefill @@ -343,21 +359,58 @@ def run_conv_qkv3(): init_state = torch.zeros(batch_size, nh, hd, hd, dtype=torch.float32, device=device) if prefill_source == "flashkda_fwd": - - def run_prefill(): - flashkda_prefill( - q, - k, - v, - raw_g, - raw_beta, - a_log, - dt_bias, - KDA_LOWER_BOUND, - init_state, - cu, + import inspect + + # vLLM 0.29.0 FlashKDA writes into caller-owned buffers taken + # from the engine workspace (kda.py:536-547 sizes them: + # out (1,T,H,D) model dtype, final_state (N,H,D,D) fp32, + # workspace uint8 of _flashkda_C.get_workspace_size(T,H,N); + # the non-checkpoint prefill call is kda.py:939-953). The + # kimi-k3 preview build returned them instead — dispatch on + # the signature so both builds run the same kernel. + if "out" in inspect.signature(flashkda_prefill).parameters: + import vllm._flashkda_C # noqa: F401 + + fk_out = torch.empty_like(q) + fk_final = torch.zeros(batch_size, nh, hd, hd, dtype=torch.float32, device=device) + fk_ws = torch.empty( + int(torch.ops._flashkda_C.get_workspace_size(nt, nh, batch_size)), + dtype=torch.uint8, + device=device, ) + def run_prefill(): + flashkda_prefill( + q, + k, + v, + raw_g, + raw_beta, + a_log, + dt_bias, + KDA_LOWER_BOUND, + init_state, + cu, + out=fk_out, + final_state=fk_final, + workspace=fk_ws, + ) + else: + + def run_prefill(): + flashkda_prefill( + q, + k, + v, + raw_g, + raw_beta, + a_log, + dt_bias, + KDA_LOWER_BOUND, + init_state, + cu, + ) + else: def run_prefill(): @@ -447,7 +500,17 @@ def run_kda_generation_benchmark( nh, hd, cw = num_v_heads, head_v_dim, d_conv dim = nh * hd - fused_ok = is_fused_kda_decode_supported(nh, hd, cw, num_spec=0, input_dtype=dtype, conv_state_dtype=dtype) + # 0.30 added the recurrent-state dtype to the decode probe as well: serving + # resolves it from get_state_dtype()[1] (vllm/models/kimi_k3/nvidia/kda.py + # :665-680@v0.30.0, resolve_kda_decode_backend -> is_fused_kda_decode_supported + # :157-178), float32 for the default mamba_ssm_cache_dtype "auto" — the dtype + # of this collector's recurrent state below. + decode_paths = _decode_paths() + fused_ok = is_fused_kda_decode_supported( + nh, hd, cw, num_spec=0, input_dtype=dtype, conv_state_dtype=dtype, recurrent_state_dtype=torch.float32 + ) + if decode_paths == {"fused"} and not fused_ok: + raise RuntimeError("AIS_KDA_DECODE_PATHS=fused, but the fused KDA decode kernel does not support this shard") conv_weight = torch.randn(3 * dim, cw, dtype=torch.float32, device=device) fused_weight = conv_weight.reshape(3, dim, cw).transpose(1, 2).contiguous() norm_weight = torch.randn(hd, dtype=torch.float32, device=device) @@ -483,7 +546,7 @@ def run_kda_generation_benchmark( "model_name": model_name, } - if fused_ok: + if fused_ok and "fused" in decode_paths: def run_fused_decode(): ops.fused_kda_decode( @@ -520,55 +583,57 @@ def run_fused_decode(): results["power_stats"], ) - def run_conv_update(): - causal_conv1d_update( - x, - conv_state, - conv_weight, - None, - activation="silu", - conv_state_indices=idx, - validate_data=True, - out=conv_out, - ) + if "packed" in decode_paths: - with benchmark_with_power( - device=device, kernel_func=run_conv_update, num_warmups=NUM_WARMUPS, num_runs=NUM_RUNS, repeat_n=1 - ) as results: - _log( - common, - results["latency_ms"], - "causal_conv1d_update", - perf_filename, - vllm_version, - device, - results["power_stats"], - ) + def run_conv_update(): + causal_conv1d_update( + x, + conv_state, + conv_weight, + None, + activation="silu", + conv_state_indices=idx, + validate_data=True, + out=conv_out, + ) - def run_packed_decode(): - fused_recurrent_kda_packed_decode( - mixed_qkv=x, - raw_g=raw_g, - raw_beta=raw_beta, - A_log=a_log, - dt_bias=dt_bias_hd, - lower_bound=KDA_LOWER_BOUND, - initial_state=state, - state_indices=idx, - ) + with benchmark_with_power( + device=device, kernel_func=run_conv_update, num_warmups=NUM_WARMUPS, num_runs=NUM_RUNS, repeat_n=1 + ) as results: + _log( + common, + results["latency_ms"], + "causal_conv1d_update", + perf_filename, + vllm_version, + device, + results["power_stats"], + ) - with benchmark_with_power( - device=device, kernel_func=run_packed_decode, num_warmups=NUM_WARMUPS, num_runs=NUM_RUNS, repeat_n=1 - ) as results: - _log( - common, - results["latency_ms"], - "fused_recurrent_kda_packed_decode", - perf_filename, - vllm_version, - device, - results["power_stats"], - ) + def run_packed_decode(): + fused_recurrent_kda_packed_decode( + mixed_qkv=x, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=a_log, + dt_bias=dt_bias_hd, + lower_bound=KDA_LOWER_BOUND, + initial_state=state, + state_indices=idx, + ) + + with benchmark_with_power( + device=device, kernel_func=run_packed_decode, num_warmups=NUM_WARMUPS, num_runs=NUM_RUNS, repeat_n=1 + ) as results: + _log( + common, + results["latency_ms"], + "fused_recurrent_kda_packed_decode", + perf_filename, + vllm_version, + device, + results["power_stats"], + ) ok += 1 except Exception as e: err += 1 @@ -823,7 +888,7 @@ def _is_glm5_next_kda(model_name: str) -> bool: # replicates; any other installed release raises KdaRuntimeNotAuditedError (a # classified failure) instead of timing a possibly different kernel path. _KDA_ARCHITECTURE_COMPAT = { - KIMI_K3_KDA_ARCHITECTURE: "vllm==0.1.dev19262", + KIMI_K3_KDA_ARCHITECTURE: "vllm==0.31.0", # 0.30.0 route -> git history; GPU smoke 2026-10-30 GLM5_NEXT_KDA_ARCHITECTURE: "vllm==0.31.0", } diff --git a/python/aisimulate/collector/vllm/collect_mhc_module.py b/python/aisimulate/collector/vllm/collect_mhc_module.py index 628bfc8331..c62a59a59b 100644 --- a/python/aisimulate/collector/vllm/collect_mhc_module.py +++ b/python/aisimulate/collector/vllm/collect_mhc_module.py @@ -28,13 +28,12 @@ # builders/selectors remain authoritative; no kernel fallback is introduced. # The campaign manifest still selects one exact release per run. # -# The file-level range spans the two audited routes; each architecture is -# additionally gated to its own audited release(s) by ``_ARCHITECTURE_COMPAT`` -# below, so DeepSeek-V4 behaviour on 0.24.0-0.25.0 is unchanged and neither -# architecture runs on an unaudited release in between (it raises instead). -# GLM-5.3-Flash moved to stock vLLM 0.31.0 (no overlay); the upper bound -# admits it while DeepSeek-V4 stays gated to its own releases. -__compat__ = "vllm>=0.24.0,<=0.31.0" +# The mhc family runs on 0.31.0 (framework_manifest.yaml family pin, +# 2026-10-30). Each architecture is additionally gated by +# ``_ARCHITECTURE_COMPAT`` below; an architecture that is not audited on the +# installed release raises instead of running. +# One release per file (owner policy 2026-09-20); earlier code is `git log -- ` away. +__compat__ = "vllm==0.31.0" # vLLM imports stay lazy in this module so that a mismatched install fails # inside collect.py's per-op error handling (after the __compat__ gate can @@ -63,8 +62,13 @@ def _mhc_tilelang_kernels(): # Per-architecture audited vLLM releases. A case whose architecture is not # audited on the installed release raises MhcRuntimeNotAuditedError (a # classified failure) instead of running an unverified dispatch. +# DeepSeek-V4 is NOT audited on 0.31.0. mhc_module_perf's architecture-free +# view (MhcKey = op_name, hc_mult, hidden_size; the DeepSeek-V4 consumer) +# would merge GLM-5.3-Flash's fused-site rows at (4, 4096) into one +# 0.31.0 table, so DeepSeek-V4 rows stay on the 0.24.0-0.30.0 tables and +# their collector route is `git log -- ` away. Re-admitting it +# needs an architecture-keyed DeepSeek-V4 lookup first (consumer change). _ARCHITECTURE_COMPAT = { - ARCHITECTURE: "vllm>=0.24.0,<=0.25.0", # Stock vllm 0.31.0 only (no overlay) — see the GLM audit below. GLM5NEXT_ARCHITECTURE: "vllm==0.31.0", } diff --git a/python/aisimulate/collector/vllm/collect_mla_bmm.py b/python/aisimulate/collector/vllm/collect_mla_bmm.py index 3218f680de..7313ddc22e 100644 --- a/python/aisimulate/collector/vllm/collect_mla_bmm.py +++ b/python/aisimulate/collector/vllm/collect_mla_bmm.py @@ -23,9 +23,12 @@ # MLA bmm call site unchanged in-image on GB300). # B200 0.25.0 qualification: both pre/post torch.bmm callbacks passed in # job 1968047. The NVIDIA bf16 BMM contract and input geometry are unchanged. -__compat__ = "vllm>=0.25.0,<=0.27.0,!=0.25.1,!=0.26.0" +__compat__ = "vllm==0.30.0" + +from importlib.metadata import ( + version as _dist_version, # setuptools/pkg_resources is absent from recent framework images +) -import pkg_resources import torch from collector.case_generator import get_mla_bmm_case_specs from collector.helper import benchmark_with_power, log_perf @@ -62,7 +65,7 @@ def _log_row(*, op_name, dtype, num_tokens, num_heads, results, perf_filename, d } ], framework="VLLM", - version=pkg_resources.get_distribution("vllm").version, + version=_dist_version("vllm"), device_name=torch.cuda.get_device_name(device), op_name=op_name, kernel_source="vllm_torch_bmm", diff --git a/python/aisimulate/collector/vllm/collect_mla_module.py b/python/aisimulate/collector/vllm/collect_mla_module.py index 82af603e83..bc249acac2 100644 --- a/python/aisimulate/collector/vllm/collect_mla_module.py +++ b/python/aisimulate/collector/vllm/collect_mla_module.py @@ -1,13 +1,16 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# +# Versions are upgraded IN PLACE (owner policy 2026-09-20): this file always +# targets the manifest pin; prior-version code is `git log -- ` +# away and prior-version DATA stays permanent under its version key. The +# 0.24->0.29 adaptations each carry a serving citation @0.29.0. On 0.29 the +# kv dtype FORKS the backend (bf16 -> FLASH_ATTN_MLA_SPARSE, fp8 -> +# FLASHMLA_SPARSE on SM90) — both keys are collected. +# # Portions adapted from vLLM dd10e03f95f94edbea1975c67ace3a35ec9a8a40. # Copyright contributors to the vLLM project. Modified for collector execution. - -# 0.27.0 audit (pod GB300 probe): DeepseekV2MLAAttention ctor params, -# _CONFIG_REGISTRY gap (glm_moe_dsa still unmapped), -# backend_supports_prefill_query_quantization (mla_attention.py) and the -# prefill selector surface are all unchanged vs the 0.24.0 citations below. -__compat__ = "vllm>=0.24.0,<=0.25.1" +__compat__ = "vllm==0.30.0" """ MLA Module Collector for vLLM — unified MLA and DSA benchmarking. @@ -26,6 +29,12 @@ config, overrides the layer-local shape fields in-memory, and then instantiates just the attention module. +Case layout: one case per (batch, heads, kv dtype, compute dtype, gemm type, +model, attention type, prefix) carrying its seq_len points; run_mla_module +builds the module once per case and measures every point on it, each with its +own KV cache and metadata (the module build was most of the per-point cost in +the 2026-10 campaign: vllm dsa_context_module alone was 59% of its GPU time). + Usage: # MLA context phase (DeepSeek-V3 style) python collect_mla_module.py --mode context --model deepseek-ai/DeepSeek-V3 @@ -47,7 +56,9 @@ import math import os import tempfile +import time import traceback +from contextlib import ExitStack from functools import partial from pathlib import Path @@ -66,6 +77,7 @@ create_common_attn_metadata, create_vllm_config, enable_engine_fused_ops, + framework_kv_block_size, setup_distributed, with_exit_stack, ) @@ -95,11 +107,15 @@ # Local model config resolution — avoid HuggingFace Hub downloads # ═══════════════════════════════════════════════════════════════════════ -# Pre-cached HF configs live in src/aisimulate_core/model_configs/ as -# "--_config.json". vLLM's ModelConfig accepts a local -# directory containing config.json, so we create a temp dir with a -# symlink when the cached file exists. -_MODEL_CONFIGS_DIR = Path(__file__).resolve().parents[2] / "src" / "aisimulate_core" / "model_configs" +# Pre-cached HF configs ship with AISim as "--_config.json". +# The directory is helper.py's constant — never recomputed here, so a +# package rename cannot silently strand this lane (the aiconfigurator -> +# aisimulate consolidation did exactly that to a hand-built copy of it). +# vLLM's ModelConfig accepts a local directory containing config.json, so +# we create a temp dir with a symlink when the cached file exists. +from collector.helper import _AIC_MODEL_CONFIG_DIR as _AIS_MODEL_CONFIG_DIR + +_MODEL_CONFIGS_DIR = Path(_AIS_MODEL_CONFIG_DIR) # Cache of model_name -> temp dir path (created once per process). _local_config_cache: dict[str, str] = {} @@ -193,6 +209,34 @@ def _generation_kv_footprint_bytes(total_tokens: int, kv_cache_dtype: str) -> in return 2 * total_tokens * _MLA_KV_ENTRY_ELEMS * bytes_per_elem +def _dsa_context_point_allowed( + batch_size: int, + seq_len: int, + prefix_len: int, + *, + new_token_budget: int, + full_seq_budget: int | None, + max_position: int | None, +) -> bool: + """Whether one DSA context point (new tokens ``seq_len`` behind ``prefix_len`` cached + tokens, ``batch_size`` requests) is in the plan: serving's chunked-prefill new-token + budget, the full-sequence budget and the model's max_position_embeddings. The plan + generator and ``run_mla_module`` (which expands a case's isl x prefix grid) share it, so + a series measures exactly the points the plan declared.""" + if batch_size * seq_len > new_token_budget: + return False + if full_seq_budget and batch_size * (seq_len + prefix_len) > full_seq_budget: + return False + return max_position is None or prefix_len + seq_len <= max_position + + +def _dsa_context_budgets(sweep=None) -> tuple[int, int | None]: + """(new-token budget, full-sequence budget) of the vLLM DSA context sweep.""" + sweep = sweep or get_mla_module_sweep_spec("vllm") + new_token_budget = getattr(sweep, "context_dsa_chunk_prefill_size", None) or sweep.context_max_tokens + return int(new_token_budget), getattr(sweep, "context_dsa_max_full_sequence_length", None) + + def get_context_test_cases(attn_type: str): """Context-phase test cases. @@ -201,16 +245,37 @@ def get_context_test_cases(attn_type: str): """ cases = [] sweep = get_mla_module_sweep_spec("vllm") + # DSA context mirrors serving's chunked prefill (matches DSV4 sparse + + # sglang DSA): cap the per-forward NEW-token budget (b*s) at the chunk + # prefill size and let the long context enter via prefix, bounded by the + # full-sequence budget b*(s+prefix). This makes rows faithful to the + # engine's query shape (isl=chunk, prefix=context) and keeps the fp8 + # sparse-decode q stride int32-safe. Falls back to the shared context + # axes when the DSA-specific budgets are not declared. + dsa_new_token_budget, dsa_full_seq_budget = _dsa_context_budgets(sweep) + dsa_prefixes = getattr(sweep, "context_dsa_prefix_lengths", None) or sweep.context_prefix_lengths for compute_dtype, kv_dtype, gemm_type in _get_precision_combos("context", attn_type): for num_heads in sweep.inner_sweep_head_counts: for b in sweep.context_batch_sizes: for s in sweep.context_sequence_lengths: - if b * s > sweep.context_max_tokens: - continue if attn_type == "dsa": - for prefix_len in sweep.context_prefix_lengths: + # new-token budget per chunk (not the shared memory cap) and the + # full-sequence budget; the per-model max_position cap is applied + # by _build_module_test_cases with the same predicate. + for prefix_len in dsa_prefixes: + if not _dsa_context_point_allowed( + b, + s, + prefix_len, + new_token_budget=dsa_new_token_budget, + full_seq_budget=dsa_full_seq_budget, + max_position=None, + ): + continue cases.append([s, b, num_heads, kv_dtype, compute_dtype, gemm_type, prefix_len]) else: + if b * s > sweep.context_max_tokens: + continue cases.append([s, b, num_heads, kv_dtype, compute_dtype, gemm_type]) return cases @@ -252,22 +317,212 @@ def get_generation_test_cases(attn_type: str): return cases +def _strided_throwaway_copy(query): + """A scratch tensor with the query's shape AND strides — never a contiguous clone. + + vLLM's own rope kernel derives ONE head stride from the query and applies it to the + key as well (csrc/libtorch_stable/pos_encoding_kernels.cu:57,68,164-165 @0.30.0, + unchanged @0.31.0). The q_pe that attention.py:534 ropes alone is a last-dim slice of + q, so its head stride is qk_head_dim (256 for GLM-5.2, 192 for DeepSeek-V3.2) while a + contiguous clone's is 64: the kernel then addresses the clone at + token*heads*64 + head*256, i.e. up to 192*(heads-1) elements PAST its end, and ropes + in place — a CUDA illegal memory access when the overrun leaves the allocation, + silent corruption of the neighbouring tensor otherwise. That was the GLM-5.2 fp8-KV + dsa_context (isl <= index_topk) illegal-address / hang family on b200, b300, gb200 + and gb300 (2026-10-14). FlashInfer's rope takes per-tensor strides (csrc/rope.cu + q_stride_h / k_stride_h @0.6.18) and is safe with either layout; matching the + query's strides is correct for both kernels. + """ + shape = tuple(query.shape) + stride = tuple(query.stride()) + if any(s == 0 for s in shape): + return query.clone() + storage_numel = 1 + sum((s - 1) * st for s, st in zip(shape, stride, strict=True)) + key = query.new_empty(storage_numel).as_strided(shape, stride) + key.copy_(query) + return key + + +def _rope_tolerating_missing_key(forward): + """Wrap a rotary module's forward so `forward(positions, query)` (key=None) works. + + vLLM 0.30's DeepseekV32 attention ropes q_pe alone in its dense-MHA FP8 branch, but + both DeepseekScalingRotaryEmbedding implementations require a key (FlashInfer op: + key.view; native: assert). Rope the query against a throwaway copy that shares the + query's strides (see _strided_throwaway_copy) and drop it. + """ + import functools + + if getattr(forward, "_aisim_rope_shim", False): + # Idempotent: run_mla_module installs the shim per case on a module that may be reused, + # and named_modules() can visit a shared rotary twice. Re-wrapping stacked one closure + # per case until "maximum recursion depth exceeded" (b200/gb200 pipelines 72387095 / + # 72387091 at e598d380: 2,286 cases each). + return forward + + @functools.wraps(forward) + def wrapped(positions, query, key=None, *args, **kwargs): + # Forward only what the caller passed: vLLM 0.30's forward_cuda signatures take + # (positions, query, key) — no `offsets` — so a positional None for offsets raised + # "takes from 3 to 4 positional arguments but 5 were given" (508 GLM-5 cases, same runs). + if key is None: + out = forward(positions, query, _strided_throwaway_copy(query), *args, **kwargs) + q = out[0] if isinstance(out, tuple) else out + return q, None + return forward(positions, query, key, *args, **kwargs) + + wrapped._aisim_rope_shim = True + return wrapped + + +def _rope_module_rejects_missing_key(module) -> bool: + """True only for rotary classes whose forward cannot take key=None. + + Serving calls ``self.rotary_emb(positions, q_pe)`` with no key + (models/deepseek_v32/attention.py:534 @0.30.0). vLLM's base RotaryEmbedding and every + subclass that reaches ops.rotary_embedding accept that (the kernel gets a null key + pointer and never touches key memory; _custom_ops.rotary_embedding key: Tensor | None), + so GLM-5.x (rope_type "default" -> RotaryEmbedding) must run exactly serving's call — + wrapping it fed vLLM's kernel a contiguous key whose head stride differs from the + query's (see _strided_throwaway_copy). Only DeepseekScalingRotaryEmbedding's own + forward_cuda rejects None (deepseek_scaling_rope.py:203-221 @0.30.0: FlashInfer op + key.view / forward_native assert key is not None); its DeepseekV4 subclass overrides + forward_cuda onto ops.rotary_embedding and accepts None again. + """ + mro = type(module).__mro__ + if not any(c.__name__ == "DeepseekScalingRotaryEmbedding" for c in mro): + return False + owner = next((c for c in mro if "forward_cuda" in vars(c)), None) + return owner is not None and owner.__name__ == "DeepseekScalingRotaryEmbedding" + + +def _install_rope_missing_key_shims(attn_module) -> tuple[list[str], list[str]]: + """Shim only the rotary submodules that reject key=None. + + Returns the class names that were wrapped and the rotary class names left on + serving's own path. + """ + patched: list[str] = [] + left: list[str] = [] + modules = attn_module.named_modules() if hasattr(attn_module, "named_modules") else () + for _name, module in modules: + if not (type(module).__name__.endswith("RotaryEmbedding") and hasattr(module, "forward")): + continue + if _rope_module_rejects_missing_key(module): + module.forward = _rope_tolerating_missing_key(module.forward) + patched.append(type(module).__name__) + else: + left.append(type(module).__name__) + return patched, left + + +def _model_max_position_embeddings(model_id: str) -> int | None: + """The model's max context length from its cached config (same file the + sglang module collector reads); None when unavailable (no cap applied).""" + config_file = os.path.join(_AIS_MODEL_CONFIG_DIR, f"{model_id.replace('/', '--')}_config.json") + if not os.path.exists(config_file): + return None + try: + with open(config_file) as f: + value = json.load(f).get("max_position_embeddings") + return int(value) if value else None + except Exception: + return None + + def _build_module_test_cases(attn_type: str, mode: str): """Build module-level test cases for a specific attention type and phase. - Output format: [seq_len, batch_size, num_heads, kv_cache_dtype, - compute_dtype, gemm_type, model_path, attn_type] + Output format: [[seq_len, ...], batch_size, num_heads, kv_cache_dtype, + compute_dtype, gemm_type, model_path, attn_type, (prefix_len)] + — one case per module series, see the grouping at the end. """ base_cases = get_context_test_cases(attn_type) if mode == "context" else get_generation_test_cases(attn_type) cases = [] + dropped_over_ceiling = 0 + dropped_non_shard_heads = 0 for model_spec in get_mla_module_model_specs(attention_type=attn_type, backend="vllm"): + max_position = _model_max_position_embeddings(model_spec.model_path) if attn_type == "dsa" else None + # Model-correlated narrowing (case_authoring.md; the MSA collectors' + # get_msa_head_counts precedent): a row that declares its real + # tensor_parallel_sizes only measures the per-rank head counts those + # shards produce. The shared inner head grid also carries 96/48/24/12 + # for the 96-head Kimi-K3 MLA module and 128 for 128-head models; on + # the 64-head GLM-5.x and the 128-head DeepSeek-V3.2 DSA rows those + # were 40-50% of the b200 dsa_context rows with no deployment behind + # them (2026-10-29). Rows without the declaration keep the full grid. + shard_heads = None + declared_tp_sizes = getattr(model_spec, "tensor_parallel_sizes", None) + if declared_tp_sizes: + shard_heads = {model_spec.native_num_heads // tp for tp in declared_tp_sizes} for base_case in base_cases: s, b, h, kv_dtype, compute_dtype, gemm_type, *rest = base_case + if shard_heads is not None and h not in shard_heads: + dropped_non_shard_heads += 1 + continue + prefix_len = rest[0] if rest else 0 + if attn_type == "dsa" and max_position is not None and prefix_len + s > max_position: + # The shared dsa_prefix_lengths axis reaches 1,048,575 (GLM-5.2's ceiling); a model + # with a smaller max_position_embeddings (DeepSeek-V3.2 163,840, GLM-5 202,752) cannot + # build an engine past it: vLLM rejects max_model_len > derived limit (7,050 cases, + # job 471043976). sglang's module collector applies the same per-model cap + # (_prefix_fits_model); mirror it here. + dropped_over_ceiling += 1 + continue case = [s, b, h, kv_dtype, compute_dtype, gemm_type, model_spec.model_path, attn_type] if rest: case.append(rest[0]) cases.append(case) - return cases + if dropped_over_ceiling or dropped_non_shard_heads: + print( + f"[vllm-mla-module-cases] {attn_type}/{mode}: dropped " + f"{dropped_over_ceiling} cases whose prefix+isl exceeds the model's max_position_embeddings and " + f"{dropped_non_shard_heads} cases whose per-rank heads no declared tensor_parallel_sizes shard produces" + ) + # One case per (batch, heads, kv, compute, gemm, model, attn_type) carrying its + # seq_len points in slot 0 and, for DSA context, its prefix points in the last slot: + # run_mla_module builds the module once per case and measures every (isl, prefix) + # point on it (largest total context first). Every point keeps its own KV cache, + # metadata and rows; only the module build is shared. The runner expands the + # isl x prefix grid through _dsa_context_point_allowed, the same predicate that + # built ``cases`` above, and the check below pins that the two agree. + if attn_type == "dsa" and mode == "context": + grid: dict[tuple, set[tuple[int, int]]] = {} + for s, b, h, kv_dtype, compute_dtype, gemm_type, model_path, attn, prefix_len in cases: + key = (b, h, kv_dtype, compute_dtype, gemm_type, model_path, attn) + grid.setdefault(key, set()).add((int(s), int(prefix_len))) + new_token_budget, full_seq_budget = _dsa_context_budgets() + grouped = [] + for key, points in grid.items(): + b, model_path = key[0], key[5] + seq_lens = sorted({s for s, _ in points}, reverse=True) + prefixes = sorted({p for _, p in points}, reverse=True) + max_position = _model_max_position_embeddings(model_path) + expanded = { + (s, p) + for s in seq_lens + for p in prefixes + if _dsa_context_point_allowed( + b, + s, + p, + new_token_budget=new_token_budget, + full_seq_budget=full_seq_budget, + max_position=max_position, + ) + } + if expanded != points: + raise RuntimeError( + f"vllm dsa context series {key}: the isl x prefix grid expands to {len(expanded)} points " + f"but the plan holds {len(points)}; _dsa_context_point_allowed and get_context_test_cases " + "disagree" + ) + grouped.append([seq_lens, *key, prefixes]) + return grouped + series: dict[tuple, list[int]] = {} + for s, *rest in cases: + series.setdefault(tuple(rest), []).append(int(s)) + return [[sorted(set(points), reverse=True), *key] for key, points in series.items()] def get_mla_context_module_test_cases(): @@ -321,19 +576,9 @@ def _create_gemm_quant_config(gemm_type: str): """ if gemm_type == "bfloat16": return None - if gemm_type in ("fp8", "fp8_block"): + if gemm_type == "fp8_block": from vllm.model_executor.layers.quantization.fp8 import Fp8Config - # Match the existing GEMM collector's fp8 contract: tensor-scaled - # weights with dynamic activations, not fp8_static. Native Fp8LinearMethod - # selects the activation granularity/kernel (fp8.py:298-322 at - # vLLM v0.25.1, commit 752a3a504485790a). - if gemm_type == "fp8": - return Fp8Config( - is_checkpoint_fp8_serialized=True, - activation_scheme="dynamic", - weight_block_size=None, - ) # vLLM requires is_checkpoint_fp8_serialized=True for block-scaled # FP8 (fp8.py raises ValueError otherwise). This routes through # Fp8LinearMethod (block_quant=True) → W8A8BlockFp8LinearOp → @@ -357,12 +602,55 @@ def _create_gemm_quant_config(gemm_type: str): def _move_module_preserving_buffers(module, device): - """Materialize meta weights without discarding initialized runtime buffers.""" + """Materialize meta weights without discarding initialized runtime buffers. + + Two construction paths reach here. The explicit-class path builds the + module on the host, so only serialized block-scaled FP8 weights are meta + and every buffer already holds its constructor value. The serving-parity + path instantiates the framework's whole model under torch.device("meta") + (_create_attention_module), so its buffers are meta as well, and each is + rebuilt from its owner's own constructor recipe: + + - rotary cos/sin tables: `rotary_emb` and `indexer_rope_emb` on + DeepseekV32Attention (models/deepseek_v32/attention.py:272-283 + @ v0.30.0), registered non-persistent by rotary_embedding/base.py:60-63 + and deepseek_scaling_rope.py:241-242 from a device-free recipe + (_compute_cos_sin_cache: torch.arange / einsum on the default device), + cast to the dtype the constructor chose (the meta tensor carries it); + - KV-quant scales `_k_scale/_v_scale/_q_scale/_prob_scale` that + MLAAttention registers through _init_kv_cache_quant -> + set_default_quant_scales(register_buffer=True) as 1.0 fp32 + (layers/attention/attention.py:129-135,184 @ v0.30.0), together with + the `_cpu` host copies the same call keeps beside them. + + Any other meta buffer is refused rather than filled: a filled table + silently changes what the kernels see (the pre-#404 collector filled RoPE + with 0.01 and these scales with 0.5). + """ if not any(parameter.is_meta for parameter in module.parameters()): return module.to(device) - buffers = dict(module.named_buffers()) - if any(buffer.is_meta for buffer in buffers.values()): - raise RuntimeError("Cannot preserve an uninitialized meta buffer in the attention module") + buffers = {} + for name, buffer in module.named_buffers(): + if buffer.is_meta: + parent_name, _, local_name = name.rpartition(".") + owner = module.get_submodule(parent_name) if parent_name else module + recompute = getattr(owner, "_compute_cos_sin_cache", None) + if local_name in ("_k_scale", "_v_scale", "_q_scale", "_prob_scale"): + buffer = torch.ones(buffer.shape, dtype=buffer.dtype) + host_name = f"{local_name}_cpu" + if getattr(getattr(owner, host_name, None), "is_meta", False): + setattr(owner, host_name, torch.ones((), dtype=torch.float32)) + elif local_name.startswith("cos_sin_cache") and callable(recompute): + table = recompute() + if table.is_meta or tuple(table.shape) != tuple(buffer.shape): + raise RuntimeError( + f"Rotary table rebuilt for {name} is {tuple(table.shape)} on {table.device}; " + f"the constructor registered {tuple(buffer.shape)}" + ) + buffer = table.to(dtype=buffer.dtype) + else: + raise RuntimeError(f"Cannot preserve an uninitialized meta buffer in the attention module: {name}") + buffers[name] = buffer module = module.to_empty(device=torch.device(device)) for name, buffer in buffers.items(): parent_name, _, local_name = name.rpartition(".") @@ -412,25 +700,35 @@ def _initialize_nvfp4_parameters(module): untouched. vLLM subsequently performs its own layout conversion/dispatch. """ from vllm import _custom_ops as ops - from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4LinearMethod - - # Native API contracts at vLLM v0.25.1, commit 752a3a504485790a: - # modelopt.py:1111-1199 stores two E2M1 values per uint8, one E4M3 scale - # per 16 inputs, and dequantizing FP32 global scales. _custom_ops.py:1492- - # 1571 accepts the inverse global scale and can return unswizzled scales; - # modelopt.py:1203-1238 performs native conversion only after loading. + from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Static + + # Native API contracts at vLLM v0.30.0: every ModelOpt linear is the + # generic ModelOptLinearMethod carrying a (weight, activation) QuantSpec + # and a CkptCtx (modelopt.py:2448-2476). The NVFP4 weight scheme + # KNvfp4Static stores two E2M1 values per uint8, one E4M3 scale per + # ctx.group_size inputs (created as a NaN sentinel that process() refuses) + # and a per-partition FP32 global scale `weight_scale_2` + # (modelopt.py:1913-1990); the W4A4 activation scheme KNvfp4Dynamic adds + # the per-partition `input_scale` (modelopt.py:1993-2028) while W4A16 has + # none. Both global scales are ModelOpt's dequantizing amax/(6*448) + # values, raw-max'ed at process_weights_after_loading. _custom_ops.py: + # 1525-1531 accepts the inverse global scale and can return unswizzled + # block scales. (v0.25.1's ModelOptNvFp4LinearMethod no longer exists.) generator = torch.Generator(device="cpu").manual_seed(0) initialized = 0 with torch.no_grad(): for name, layer in module.named_modules(): - if not isinstance(getattr(layer, "quant_method", None), ModelOptNvFp4LinearMethod): + quant_method = getattr(layer, "quant_method", None) + spec = getattr(quant_method, "spec", None) + if spec is None or getattr(spec, "weight", None) != kNvfp4Static: continue weight, scales = layer.weight, layer.weight_scale if weight.dtype != torch.uint8 or scales.dtype != torch.float8_e4m3fn: raise ValueError(f"Unexpected native NVFP4 storage for {name}") if weight.ndim != 2 or tuple(scales.shape) != (weight.shape[0], weight.shape[1] // 8): raise ValueError(f"Unexpected native NVFP4 block layout for {name}") - if layer.quant_method.quant_config.group_size != 16 or weight.shape[1] % 8: + group_size = getattr(getattr(quant_method, "ctx", None), "group_size", None) + if group_size != 16 or weight.shape[1] % 8: raise ValueError(f"Native NVFP4 quantization requires 16-input blocks for {name}") input_size = weight.shape[1] * 2 rows_per_chunk = max(1, 1048576 // input_size) @@ -439,7 +737,9 @@ def _initialize_nvfp4_parameters(module): # These fixed source ranges are unrelated to timing/accuracy targets. weight_global_scale = 0.08 / (6.0 * 448.0) layer.weight_scale_2.fill_(weight_global_scale) - layer.input_scale.fill_(1.0 / (6.0 * 448.0)) + input_scale = getattr(layer, "input_scale", None) # W4A4 only (KNvfp4Dynamic) + if input_scale is not None: + input_scale.fill_(1.0 / (6.0 * 448.0)) inverse_scale = torch.tensor(1.0 / weight_global_scale, dtype=torch.float32, device=weight.device) for start in range(0, weight.shape[0], rows_per_chunk): end = min(start + rows_per_chunk, weight.shape[0]) @@ -473,6 +773,7 @@ def _create_attention_module( gemm_type: str = "bfloat16", device: str = "cuda:0", is_context: bool = True, + max_model_len: int | None = None, ): """ Create a DeepseekV2MLAAttention module from vLLM's own modeling code. @@ -486,6 +787,9 @@ def _create_attention_module( model_path: HuggingFace model path (e.g. "deepseek-ai/DeepSeek-V3.2"). attn_type: Attention type ("mla" or "dsa"). max_seq_len: Total context length, including the cached prefix. + max_model_len: The engine's max_model_len (default: max_seq_len, at least 4096). + Indexer buffers scale with it, so a series passes the model's + max_position to time every point at the size serving runs. max_query_len: New tokens per request for a context forward pass. use_prefill_fp8: When True, opt in to FP8 prefill query compute via ``attention_config.use_prefill_query_quantization`` plus an @@ -500,21 +804,32 @@ def _create_attention_module( _create_kv_cache_and_metadata fails closed if the decision does not match the case label. gemm_type: Precision for linear-layer GEMMs — "bfloat16", - "fp8", "fp8_block", or "nvfp4". + "fp8_block", or "nvfp4". """ from vllm.model_executor.models.deepseek_v2 import DeepseekV2MLAAttention local_model_path = _resolve_model_path(model_path) - block_size = 64 + # The module is built at the framework default page; the real page is asked from the + # backend the module selected (_create_kv_cache_and_metadata, utils.framework_kv_block_size). + block_size = None # unconstrained selection; _create_kv_cache_and_metadata asks the selected backend for the page # seq_len includes the current token; generation caches only seq_len - 1. # Keep exact-limit models such as Kimi-K2-Instruct at their declared limit. - max_model_len = max(max_seq_len, 4096) + max_model_len = max(max_seq_len, max_model_len or 4096) num_kv_cache_blocks = max( - 1 + math.ceil((max_seq_len + 1) / block_size) * max_batch_size, + 1 + + math.ceil((max_seq_len + 1) / 16) + * max_batch_size, # upper bound at the smallest page; recomputed per page below 8192, ) + # FIXME(gap): the module is built at TP=1 with the per-rank head count, so a check that + # multiplies by the TP world size sees the shard. vLLM 0.30's sparse-MLA masked-MHA prefill + # (SM100, FA4, BF16 KV) admits only model-level head counts 128/64 + # (sparse_mla_attention.py:93-127, called with num_heads * tp at :708), and _use_masked_mha + # reads its thresholds by tensor_parallel_size (mla_attention.py:1844). So for a TP>1 shard, + # e.g. GLM-5.2 TP8 BF16 KV, serving runs masked MHA for pure prefills up to 64K while this + # module times sparse MQA. Affects SM100 BF16-KV DSA context rows below the full head count. # Determine kv cache dtype string for sparse MLA. # For DSA (DeepSeekV3.2), fp8 uses the custom ``fp8_ds_mla`` 656-byte # cache format (512B quantized NoPE + 16B scales + 128B RoPE). @@ -626,26 +941,62 @@ def _create_attention_module( from vllm.utils.torch_utils import set_default_torch_dtype with set_current_vllm_config(vllm_config), set_default_torch_dtype(vllm_config.model_config.dtype): - attn_module = DeepseekV2MLAAttention( - vllm_config=vllm_config, - config=hf_config, - hidden_size=hf_config.hidden_size, - num_heads=num_heads, - qk_nope_head_dim=hf_config.qk_nope_head_dim, - qk_rope_head_dim=hf_config.qk_rope_head_dim, - v_head_dim=hf_config.v_head_dim, - q_lora_rank=hf_config.q_lora_rank if hasattr(hf_config, "q_lora_rank") else None, - kv_lora_rank=hf_config.kv_lora_rank, - max_position_embeddings=hf_config.max_position_embeddings, - cache_config=vllm_config.cache_config, - quant_config=vllm_config.quant_config, - prefix="model.layers.0.self_attn", - topk_indices_buffer=topk_indices_buffer, - ) + if attn_type == "dsa": + # GENERIC serving-parity construction (owner rule 2026-09-19: + # config in -> the module SERVING would build comes out; never a + # per-model class import — that is how the prelude-fusion drift + # was born). Resolve the model class through the framework's own + # registry (exactly serving's routing, registry.resolve_model_cls), + # instantiate the WHOLE model on the meta device (free), and + # extract layer 0's attention submodule. Attention layers register + # themselves into static_forward_context during init, which is + # what _create_kv_cache_and_metadata reads — same as serving. + from vllm.model_executor.models.registry import ModelRegistry + + model_cls, _resolved = ModelRegistry.resolve_model_cls(hf_config.architectures, vllm_config.model_config) + with torch.device("meta"): + _meta_model = model_cls(vllm_config=vllm_config, prefix="") + attn_module = None + for _n, _m in _meta_model.named_modules(): + if _n.endswith("layers.0.self_attn"): + attn_module = _m + break + if attn_module is None: # no fallback: raise, never substitute + raise RuntimeError( + f"framework model {model_cls.__name__} exposes no " + "layers.0.self_attn — extend the extraction, do not pin a class" + ) + # the model allocated its shared topk buffer on meta; swap in the + # real one the collector allocated (same shape contract) + for _n, _m in attn_module.named_modules(): + if getattr(_m, "topk_indices_buffer", None) is not None: + _m.topk_indices_buffer = topk_indices_buffer + if getattr(attn_module, "topk_indices_buffer", None) is not None: + attn_module.topk_indices_buffer = topk_indices_buffer + else: + attn_module = DeepseekV2MLAAttention( + vllm_config=vllm_config, + config=hf_config, + hidden_size=hf_config.hidden_size, + num_heads=num_heads, + qk_nope_head_dim=hf_config.qk_nope_head_dim, + qk_rope_head_dim=hf_config.qk_rope_head_dim, + v_head_dim=hf_config.v_head_dim, + q_lora_rank=hf_config.q_lora_rank if hasattr(hf_config, "q_lora_rank") else None, + kv_lora_rank=hf_config.kv_lora_rank, + max_position_embeddings=hf_config.max_position_embeddings, + cache_config=vllm_config.cache_config, + quant_config=vllm_config.quant_config, + prefix="model.layers.0.self_attn", + topk_indices_buffer=topk_indices_buffer, + ) # Serialized block-scaled FP8 creates weight params on meta device; # to_empty() also discards already-initialized buffers. Preserve their - # constructor values, including nonpersistent RoPE tables, across it. + # constructor values, including nonpersistent RoPE tables, across it. On + # the serving-parity path above the whole module is meta, so the RoPE + # tables and KV-quant scales are rebuilt from their owners' constructor + # recipe instead of preserved. attn_module = _move_module_preserving_buffers(attn_module, device) attn_module.eval() attn_module.requires_grad_(False) @@ -753,6 +1104,13 @@ def _populate_indexer_kv_cache( # ═══════════════════════════════════════════════════════════════════════ +# Per-case phase timings. run_mla_module prints one "[... module timing]" line +# per case so a run log shows where the module collectors spend their time +# (framework setup per case vs. the timed forwards); _create_kv_cache_and_metadata +# records the indexer-cache fill here because its return tuple is a contract. +_PHASE_TIMINGS: dict[str, float] = {} + + def _create_kv_cache_and_metadata( vllm_config, attn_type: str, @@ -765,10 +1123,24 @@ def _create_kv_cache_and_metadata( indexer=None, ): """Create KV cache and attention metadata for benchmarking.""" + _PHASE_TIMINGS["indexer_fill"] = 0.0 hf_config = vllm_config.model_config.hf_text_config kv_lora_rank = hf_config.kv_lora_rank qk_rope_head_dim = hf_config.qk_rope_head_dim - block_size = vllm_config.cache_config.block_size + # page = what the framework would run the SELECTED backend with. Serving sizes the cache + # block from the FIRST non-SSM attention layer (platforms/interface.py:588-640 @0.30.0). + # For DSA that is the indexer cache — deepseek_v2.py:1175 builds the Indexer (its + # k_cache registers at :724) before :1211 registers .attn — whose CUDA backend accepts + # only 64 (v1/attention/backends/mla/indexer.py:203-204), so FLASHINFER_MLA_SPARSE + # ([32, 64]) runs at 64, not at its own preferred 32. + from vllm.platforms import current_platform + + _attn_layer = vllm_config.compilation_config.static_forward_context["model.layers.0.self_attn.attn"] + block_size = framework_kv_block_size( + _attn_layer.get_attn_backend(), + cache_backend_cls=current_platform._find_non_ssm_backend(vllm_config), + ) + vllm_config.cache_config.block_size = block_size is_dsa = attn_type == "dsa" if is_dsa and indexer is None: @@ -889,17 +1261,67 @@ def _create_kv_cache_and_metadata( device=device, ) indexer_builder_cls = indexer_layer.get_attn_backend().get_builder_cls() - indexer_builder = indexer_builder_cls(indexer_spec, [indexer_layer_name], vllm_config, torch.device(device)) + # 0.29.0: builders may demand block_table_width (indexer.py:519,531); + # mirror serving's population exactly (v1/worker/utils.py:276-283) + # using vllm's own width helper (v1/worker/block_table.py:29-49). + builder_kwargs = {} + if getattr(indexer_builder_cls, "requires_block_table_width", False): + from vllm.v1.worker.block_table import get_block_table_width + + max_num_blocks = indexer_spec.max_num_blocks_per_req(vllm_config, vllm_config.model_config.max_model_len) + builder_kwargs["block_table_width"] = get_block_table_width(max_num_blocks, indexer_spec.block_size) + # Serving allocates the block table AT block_table_width columns + # (v1/worker/block_table.py), so builders may assume + # table.shape[1] == width. The family-100 decode path acts on + # that assumption: it copies the metadata's table into + # expanded_block_table_buffer sized by the width + # (v1/attention/backends/mla/indexer.py:609,745@0.29.0) and + # crashes on our ceil(seq/block)-column table ("expanded size + # (64) must match existing size (32)"). Zero-pad to the width — + # padding blocks are never dereferenced for absent tokens. SM90 + # paths never read the expanded buffer, which is why this stayed + # latent on H20 (found on B300, 2026-09-20). + _width = builder_kwargs["block_table_width"] + _bt = common_attn_metadata.block_table_tensor + if _bt.shape[1] < _width: + _pad = torch.zeros( + (_bt.shape[0], _width - _bt.shape[1]), + dtype=_bt.dtype, + device=_bt.device, + ) + common_attn_metadata.block_table_tensor = torch.cat([_bt, _pad], dim=1) + indexer_builder = indexer_builder_cls( + indexer_spec, [indexer_layer_name], vllm_config, torch.device(device), **builder_kwargs + ) indexer_metadata = indexer_builder.build( common_prefix_len=prefix_len, common_attn_metadata=common_attn_metadata, ) + if os.environ.get("AIS_DEBUG_IDXMETA") or os.environ.get( + "AIC_DEBUG_IDXMETA" + ): # diagnostic dump, no behavior change + _pre = getattr(indexer_metadata, "prefill", None) + print( + f"[idxmeta] type={type(indexer_metadata).__name__} " + f"prefill={type(_pre).__name__ if _pre is not None else None} " + f"decode={getattr(indexer_metadata, 'decode', None) is not None}" + ) + if _pre is not None: + _ch = getattr(_pre, "chunks", None) + print(f"[idxmeta] chunks={None if _ch is None else len(_ch)}") + for _c in (_ch or [])[:1]: + for _f in dir(_c): + _v = getattr(_c, _f, None) + if not _f.startswith("_") and isinstance(_v, (int, bool, float)): + print(f"[idxmeta] chunk.{_f} = {_v}") + _fill_start = time.perf_counter() _populate_indexer_kv_cache( indexer_kv_cache=indexer_kv_cache, common_attn_metadata=common_attn_metadata, context_lens=[tensor.shape[0] for tensor in kv_c_contexts], indexer=indexer, ) + _PHASE_TIMINGS["indexer_fill"] = time.perf_counter() - _fill_start return kv_cache, attn_metadata, common_attn_metadata, indexer_kv_cache, indexer_metadata @@ -909,17 +1331,85 @@ def _create_kv_cache_and_metadata( # ═══════════════════════════════════════════════════════════════════════ +def _series_seq_lens(seq_len) -> list[int]: + """One point (int) or the series of points (list) a case measures on one module. + + Series points run largest first: the module and every per-point buffer are + sized by the largest point, so a capacity failure shows up on the first + point instead of after the smaller ones were measured. + """ + if isinstance(seq_len, bool): + raise TypeError("seq_len must be an int or a list of ints") + if isinstance(seq_len, int): + return [int(seq_len)] + seq_lens = sorted({int(value) for value in seq_len}, reverse=True) + if not seq_lens: + raise ValueError("a module series needs at least one seq_len") + return seq_lens + + +def _series_points( + seq_lens: list[int], + prefix_len, + *, + batch_size: int, + model_path: str, +) -> list[tuple[int, int]]: + """The (seq_len, prefix_len) points a case measures on one module, largest total + context first (a capacity failure shows up on the first point). + + ``prefix_len`` is one prefix (int: every seq_len behind it) or the case's prefix + list (the plan's isl x prefix grid, expanded through _dsa_context_point_allowed + exactly as the plan generator filtered it).""" + if isinstance(prefix_len, bool): + raise TypeError("prefix_len must be an int or a list of ints") + if isinstance(prefix_len, int): + points = [(s, int(prefix_len)) for s in seq_lens] + else: + prefixes = sorted({int(value) for value in prefix_len}, reverse=True) + if not prefixes: + raise ValueError("a module series needs at least one prefix_len") + new_token_budget, full_seq_budget = _dsa_context_budgets() + max_position = _model_max_position_embeddings(model_path) + points = [ + (s, p) + for s in seq_lens + for p in prefixes + if _dsa_context_point_allowed( + batch_size, + s, + p, + new_token_budget=new_token_budget, + full_seq_budget=full_seq_budget, + max_position=max_position, + ) + ] + if not points: + raise ValueError("the module series' isl x prefix grid has no point inside the plan budgets") + return sorted(points, key=lambda sp: (sp[0] + sp[1], sp[0]), reverse=True) + + +def _is_fatal_cuda_error(exc: BaseException) -> bool: + """A sticky CUDA failure: nothing measured on this context afterwards is trustworthy.""" + accelerator_error = getattr(torch, "AcceleratorError", None) + if accelerator_error is not None and isinstance(exc, accelerator_error): + return True + text = str(exc).lower() + markers = ("cuda error", "illegal memory access", "cudaerrorillegaladdress", "device-side assert") + return any(marker in text for marker in markers) + + @with_exit_stack def run_mla_module( exit_stack, - seq_len: int, + seq_len: int | list[int], batch_size: int, num_heads: int, kv_cache_dtype: str, compute_dtype: str, gemm_type: str, perf_filename: str, - prefix_len: int = 0, + prefix_len: int | list[int] = 0, *, model_path: str, attn_type: str, @@ -927,7 +1417,23 @@ def run_mla_module( warming_up: int = 10, test_ite: int = 6, ): - """Run a single MLA / DSA module-level benchmark point.""" + """Run one MLA / DSA module series: one module, one measured point per (``seq_len``, prefix). + + ``seq_len`` is the query length (context) or the KV length (generation) of a + point, or the list of them that share (batch, heads, dtypes, gemm, model); + ``prefix_len`` is one cached-prefix length or, for a DSA context series, the + case's prefix list (the points are the isl x prefix grid inside the plan + budgets, see _series_points). The module (vLLM config, synthetic/quantized + weights, rope shims) is built once for the largest point; the KV cache, the indexer cache and the + attention metadata are still built per point exactly as a single-point run + does, so every row keeps its serving-parity construction. A point that fails + for a non-fatal reason (OOM, a kernel limit) is recorded and the remaining + points are still measured; the failure is raised after the series so + collect.py classifies the case (the point's own exception when it is the + only one). A sticky CUDA failure aborts the series at once. + + Returns the full-indexer latencies of the measured points, largest first. + """ if attn_type not in {"mla", "dsa"}: raise ValueError(f"unsupported vLLM attention type: {attn_type!r}") if kv_cache_dtype not in {"bfloat16", "fp8"}: @@ -940,6 +1446,16 @@ def run_mla_module( "(determine_prefill_query_data_type, mla_attention.py:1462-1466 " f"@0.24.0); got kv_cache_dtype={kv_cache_dtype!r}" ) + seq_lens = _series_seq_lens(seq_len) + + _t_series = _t_last = time.perf_counter() + _setup_phase: dict[str, float] = {} + + def _lap(phases: dict[str, float], name: str) -> None: + nonlocal _t_last + now = time.perf_counter() + phases[name] = now - _t_last + _t_last = now setup_distributed(device) torch.cuda.set_device(device) @@ -957,35 +1473,117 @@ def run_mla_module( use_fp8_kv_cache = kv_cache_dtype == "fp8" use_prefill_fp8 = compute_dtype == "fp8" is_context = "context" in perf_filename - prefix_len = int(prefix_len) if is_context else 0 + points = _series_points(seq_lens, prefix_len if is_context else 0, batch_size=batch_size, model_path=model_path) + prefix_lens = sorted({p for _, p in points}, reverse=True) + prefix_desc = prefix_lens[0] if len(prefix_lens) == 1 else prefix_lens phase = "context" if is_context else "generation" variant = attn_type.upper() + max_seq_len = max(s for s, _ in points) + max_total_len = max(s + p for s, p in points) print( - f"\n[{variant} module] {phase} b={batch_size}, s={seq_len}, " - f"prefix={prefix_len}, heads={num_heads}, gemm={gemm_type}, " + f"\n[{variant} module] {phase} b={batch_size}, s={seq_lens if len(seq_lens) > 1 else seq_lens[0]}, " + f"prefix={prefix_desc}, heads={num_heads}, gemm={gemm_type}, " f"compute={compute_dtype}, kv={kv_cache_dtype}, model={model_path}" ) - # 1. Create attention module + # 1. Create attention module, sized for the series' largest point and at the + # model's max_position_embeddings as max_model_len — serving's default, and + # independent of the batch: the DSA indexer's decode logits + # [rows, max_model_len] (sparse_attn_indexer.py:711-721 @0.30.0) and prefill + # buffer max_model_len * 40 (mla/indexer.py:722-731) scale with it, so a + # series-max size would bend batch scaling within one table. The other sizes + # only reach capacity declarations (max_num_batched_tokens, the top-k index + # buffer); the backend selection is dtype/SM driven and every point's KV + # cache + metadata is built below for that point's own shape. attn_module, vllm_config = _create_attention_module( model_path=model_path, attn_type=attn_type, num_heads=num_heads, use_fp8_kv_cache=use_fp8_kv_cache, use_prefill_fp8=use_prefill_fp8, - max_seq_len=prefix_len + seq_len, + max_seq_len=max_total_len, max_batch_size=batch_size, - max_query_len=seq_len, + max_query_len=max_seq_len, gemm_type=gemm_type, device=device, is_context=is_context, + max_model_len=_model_max_position_embeddings(model_path), ) + _lap(_setup_phase, "module") + if attn_type == "dsa" and use_fp8_kv_cache: + # vLLM 0.30 DeepseekV32 attention: the short-prefill dense-MHA branch with an FP8 query + # calls `self.rotary_emb(positions, q_pe)[0]` (models/deepseek_v32/attention.py:534) — + # key=None. Serving takes that branch on the SM100 family with an FP8 KV cache + # (_fp8_query = fp8 KV and FlashInferMLASparseImpl.supports_quant_query_input, + # attention.py:239 / flashinfer_mla_sparse.py:481; _use_sparse_mha for prompts <= + # index_topk). The yarn rope of DeepSeek-V3.2 (DeepseekScalingRotaryEmbedding) rejects + # key=None — FlashInfer op `key.view(...)` ("'NoneType' object has no attribute 'view'", + # 3.9k b200 + ~3k b300 + ~3k gb200 cases) / forward_native `assert key is not None` + # (deepseek_scaling_rope.py:123; the ed9109ec attempt to run native moved every case + # onto that assert: gb200 pipeline 72343556) — so for THAT class the call's intent, + # "rope q_pe only", is given to the module: rope the query against a throwaway copy + # and return (query, None). GLM-5.x's default rope (RotaryEmbedding) accepts key=None + # natively and is left on serving's exact call: wrapping it too handed vLLM's own rope + # kernel a contiguous key with the wrong head stride -> out-of-bounds in-place writes, + # the GLM-5.2 fp8-KV illegal-address / stall family (2026-10-14, see + # _strided_throwaway_copy). Kernels and labels unchanged either way. Drop the shim once + # vLLM's :534 call or DeepseekScalingRotaryEmbedding accepts key=None. + _patched, _left = _install_rope_missing_key_shims(attn_module) + if _patched: + print( + f"[vllm-mla-module] dsa fp8: rope shim on {sorted(set(_patched))} (reject key=None; " + f"vLLM 0.30 attention.py:534 ropes q_pe alone); on serving's own call: {sorted(set(_left))}" + ) # 1b. Process weights (FP8 quantization + create W_UK_T / W_UV for MLA) _process_module_weights(attn_module, vllm_config, device) + _lap(_setup_phase, "weights") - # 2. Create KV cache + metadata - with set_current_vllm_config(vllm_config): + # set_current_vllm_config — needed by quantised layers and RoPE for every + # forward of the series; entered once, after _cleanup was registered, so it + # unwinds before the workspace is released. + exit_stack.enter_context(set_current_vllm_config(vllm_config)) + + attn_layer_name = "model.layers.0.self_attn.attn" + indexer_layer_name = "model.layers.0.self_attn.indexer.k_cache" + forward_ctx = vllm_config.compilation_config.static_forward_context + hidden_size = vllm_config.model_config.hf_text_config.hidden_size + + # Index reuse (GLM-5.x index_topk_freq / index_topk_pattern): the full dry run below + # populates the shared top-k buffer; serving's reuse layers run with skip_topk=True, + # keeping the MLA projections and sparse attention while skipping the indexer + # (vllm/models/deepseek_v32/attention.py:169-177,226,297,324,406,460 @v0.30.0; the + # generic DeepseekV2 wrapper keeps the flag on mla_attn, layers/mla.py:109,147,232). + # Emit both op_name variants for checkpoints declaring reuse; never a model-averaged + # latency in a full-indexer row. + hf_text_config = vllm_config.model_config.hf_text_config + reuse_pattern = getattr(hf_text_config, "index_topk_pattern", None) + has_reuse = attn_type == "dsa" and ( + "S" in reuse_pattern if reuse_pattern is not None else (getattr(hf_text_config, "index_topk_freq", 1) or 1) > 1 + ) + # 0.30 DeepseekV32Attention IS the MLA layer and owns skip_topk; the legacy generic + # DeepseekV2MLAAttention nests it under mla_attn. + skip_owner = next( + (c for c in (attn_module, getattr(attn_module, "mla_attn", None)) if c is not None and hasattr(c, "skip_topk")), + None, + ) + if has_reuse: + if skip_owner is None: + raise RuntimeError("vLLM DSA reuse requires the native skip_topk execution path") + exit_stack.callback(setattr, skip_owner, "skip_topk", skip_owner.skip_topk) + + # Record architecture to distinguish different DSA models in the perf CSV. + # perf_database uses this as a dict key when loading data. + # Aligns with sdk/models.py which uses architectures[0] throughout. + hf_cfg = vllm_config.model_config.hf_config + architecture = getattr(hf_cfg, "architectures", [getattr(hf_cfg, "model_type", "unknown")])[0] + # legacy generic module nests the layer (module.mla_attn.mla_attn); + # DeepseekV32Attention IS the MLAAttention layer itself + mla_layer = attn_module if not hasattr(attn_module, "mla_attn") else attn_module.mla_attn.mla_attn + actual_kv_cache_dtype = "fp8" if mla_layer.kv_cache_dtype.startswith("fp8") else "bfloat16" + + def _run_point(point_stack, seq_len: int, prefix_len: int, point_phase: dict[str, float]) -> float: + # 2. Create KV cache + metadata for this point kv_cache, attn_metadata, common_attn_metadata, indexer_kv_cache, indexer_metadata = ( _create_kv_cache_and_metadata( vllm_config=vllm_config, @@ -999,226 +1597,260 @@ def run_mla_module( indexer=getattr(attn_module, "indexer", None), ) ) + _lap(point_phase, "kv_metadata") - # 2b. Bind KV cache to the 0.24.0 attention layer. - attn_layer_name = "model.layers.0.self_attn.attn" - forward_ctx = vllm_config.compilation_config.static_forward_context - forward_ctx[attn_layer_name].kv_cache = kv_cache + # 2b. Bind KV cache to the 0.24.0 attention layer. + forward_ctx[attn_layer_name].kv_cache = kv_cache - # For DSA, also bind the indexer's KV cache. - indexer_layer_name = "model.layers.0.self_attn.indexer.k_cache" - if indexer_kv_cache is not None and indexer_layer_name in forward_ctx: - forward_ctx[indexer_layer_name].kv_cache = indexer_kv_cache + # For DSA, also bind the indexer's KV cache. + if indexer_kv_cache is not None and indexer_layer_name in forward_ctx: + forward_ctx[indexer_layer_name].kv_cache = indexer_kv_cache - # 3. Input tensors - hidden_size = vllm_config.model_config.hf_text_config.hidden_size - if is_context: - num_tokens = seq_len * batch_size - positions = ( - torch.arange(prefix_len, prefix_len + seq_len, device=device, dtype=torch.long) - .unsqueeze(0) - .expand(batch_size, -1) - .reshape(-1) - .contiguous() - ) - else: - num_tokens = batch_size - positions = torch.full( - (batch_size,), - seq_len - 1, + # 3. Input tensors + if is_context: + num_tokens = seq_len * batch_size + positions = ( + torch.arange(prefix_len, prefix_len + seq_len, device=device, dtype=torch.long) + .unsqueeze(0) + .expand(batch_size, -1) + .reshape(-1) + .contiguous() + ) + else: + num_tokens = batch_size + positions = torch.full( + (batch_size,), + seq_len - 1, + device=device, + dtype=torch.long, + ) + + hidden_states = torch.full( + (num_tokens, hidden_size), + 0.01, + dtype=torch.bfloat16, device=device, - dtype=torch.long, ) - hidden_states = torch.full( - (num_tokens, hidden_size), - 0.01, - dtype=torch.bfloat16, - device=device, - ) - - # 4. Dry run - # set_current_vllm_config — needed by quantised layers and RoPE. - # set_forward_context — provides attn_metadata + kv_cache to the - # MLAAttention.forward() path (it calls get_forward_context()). - # FIXME(kernel-limit): on SM100, FlashInfer's trtllm-gen sparse-MLA decode - # ships kernels only for tileSizeQ >= 8. DSA cases whose per-rank head - # count is 1/2/4 (tileSizeQ = heads for q_len 1) raise "Missing TRTLLM-GEN - # kernel (decode)" (flashinfer csrc/trtllm_fmha_kernel_launcher.cu:272, - # flashinfer 0.6.12) whenever vLLM selects FLASHINFER_MLA_SPARSE — always - # for FP8 KV, and for BF16 KV at heads <= 16 (platforms/cuda.py:98-116 - # @0.24.0). Heads 8..128 pass; boundary measured on B200. Serving fails - # identically, so the affected cases stay observed runtime failures. - # Re-verify on the next vLLM/FlashInfer bump. - # FIXME(kernel-limit): on SM120 and SM89, vLLM's dense-MLA decode backend - # is TRITON_MLA (SM120: platforms/cuda.py:130-134; SM89: the else-branch - # priority list at cuda.py:135-142, where FLASH_ATTN_MLA/FLASHMLA/ - # FLASHINFER_MLA all reject major != 9/10 and TRITON_MLA is the only - # eligible backend @0.24.0). Two measured limits (RTX PRO 6000 - # Blackwell; serving fails identically; re-verify on the next vLLM - # bump): - # 1. FP8 KV + q-heads >= 2 routes to the grouped decode kernel, which - # requests 102400B shared memory vs SM120's 101376B limit -> - # OutOfResources (heads == 1 takes the non-grouped kernel and - # passes; bf16 KV passes). vLLM's overflow guard only drops - # num_stages at BLOCK_DMODEL >= 1024, which the MLA Lk=576 path - # (BLOCK_DMODEL=512) never reaches (triton_decode_attention.py - # :490-532 @0.24.0). Upstream fix in flight: vllm#46728 (open PR, - # num_stages=1 fallback for exactly this tile, fixes vllm#46721). - # 2. Decode batches whose total cached tokens exceed ~2^31/576 raise - # a deterministic illegal memory access (reproduced in isolation - # on a clean GPU): largest passing batch*seq = 2.10M tokens, - # smallest failing = 4.19M, bracketing 2^31 / 576 elements = 3.73M - # — consistent with int32 offset overflow in the Triton decode - # kernel's KV indexing. Affects both KV dtypes (fp8 only at - # heads == 1, where limit 1 does not fire first). Same family on - # SM89 (L40S): 33/400 sampled generation cases, smallest failing - # batch*seq again 4.19M tokens (bf16 KV; fp8-KV combos are not - # declared below SM90). - # FIXME(kernel-limit): on SM120, every DSA case fails: the CUDA sparse - # attention indexer hard-requires DeepGEMM (sparse_attn_indexer.py:468-472 - # @0.24.0), whose fp8 MQA-logits kernels ship for SM90/SM100 only - # (vllm/third_party/deep_gemm .../impls/sm{90,100}_fp8_mqa_logits.cuh) and - # assert "Unsupported architecture" (deepgemm csrc/apis/attention.hpp:184); - # cases with fp8_block linears die even earlier in DeepGEMM's scale-factor - # layout transform ("Unknown SF transformation", layout.hpp:59). Serving - # fails identically. Upstream: vllm#45317 (gap report), TRITON_MLA_SPARSE - # backend for SM8x/11x/12x in vllm#38476/#47629 (open PRs), sparse decode - # in vllm#47527; DeepGEMM SM120 support in DeepGEMM#318 (open PR). - # Re-verify on the next vLLM/DeepGEMM bump. - exit_stack.enter_context(set_current_vllm_config(vllm_config)) - attn_metadata_dict = {attn_layer_name: attn_metadata} - if indexer_metadata is not None: - attn_metadata_dict[indexer_layer_name] = indexer_metadata - # Native MLA cache insertion reads the per-layer forward-context mapping, - # separately from the attention metadata. Omitting it silently suppresses - # current-token writes (vLLM mla_attention.py:1037-1058 and - # attention.py:726-766 @ vLLM v0.25.1, commit 752a3a504485790a). - slot_mapping_dict = {attn_layer_name: common_attn_metadata.slot_mapping} - if indexer_metadata is not None: - slot_mapping_dict[indexer_layer_name] = common_attn_metadata.slot_mapping - exit_stack.enter_context(set_forward_context(attn_metadata_dict, vllm_config, slot_mapping=slot_mapping_dict)) - try: - with torch.inference_mode(): - attn_module.forward(positions, hidden_states, None) - except torch.cuda.OutOfMemoryError as e: - print(f" Dry run OOM: {e}") - # Let collect.py record the capacity failure. Returning normally would - # mark a task done even though it emitted no performance row. - raise - except Exception as e: - print(f" Dry run failed: {e}") - traceback.print_exc() - # Propagate to collect.py's worker so the failure is recorded in the - # error queue. Swallowing here lets a whole op complete with zero - # rows while the pipeline reports success. - raise - - # The full dry run above populates the shared top-k indices. Native vLLM - # uses skip_topk to reuse that buffer while retaining MLA projections and - # sparse attention (deepseek_v2.py:1080-1104,1150-1162 and layers/mla.py - # MultiHeadLatentAttentionWrapper @ vLLM v0.25.1, commit 752a3a504485790a). - # Emit both existing op_name variants for checkpoints declaring reuse; - # no model-averaged latency is written into a full-indexer row. - hf_text_config = vllm_config.model_config.hf_text_config - reuse_pattern = getattr(hf_text_config, "index_topk_pattern", None) - has_reuse = attn_type == "dsa" and ( - "S" in reuse_pattern if reuse_pattern is not None else (getattr(hf_text_config, "index_topk_freq", 1) or 1) > 1 - ) - if has_reuse: - if not hasattr(attn_module.mla_attn, "skip_topk"): - raise RuntimeError("vLLM DSA reuse requires the native skip_topk execution path") - exit_stack.callback(setattr, attn_module.mla_attn, "skip_topk", attn_module.mla_attn.skip_topk) - - def measure_and_log(skip_indexer): - if has_reuse: - attn_module.mla_attn.skip_topk = skip_indexer + # 4. Dry run + # set_forward_context — provides attn_metadata + kv_cache to the + # MLAAttention.forward() path (it calls get_forward_context()). + # FIXME(kernel-limit): on SM100, FlashInfer's trtllm-gen sparse-MLA decode + # ships kernels only for tileSizeQ >= 8. DSA cases whose per-rank head + # count is 1/2/4 (tileSizeQ = heads for q_len 1) raise "Missing TRTLLM-GEN + # kernel (decode)" (flashinfer csrc/trtllm_fmha_kernel_launcher.cu:272, + # flashinfer 0.6.12) whenever vLLM selects FLASHINFER_MLA_SPARSE — always + # for FP8 KV, and for BF16 KV at heads <= 16 (platforms/cuda.py:98-116 + # @0.24.0). Heads 8..128 pass; boundary measured on B200. Serving fails + # identically; run_mla_module_worker raises for them before the build. + # Re-verify on the next vLLM/FlashInfer bump. + # FIXME(kernel-limit): on SM120 and SM89, vLLM's dense-MLA decode backend + # is TRITON_MLA (SM120: platforms/cuda.py:130-134; SM89: the else-branch + # priority list at cuda.py:135-142, where FLASH_ATTN_MLA/FLASHMLA/ + # FLASHINFER_MLA all reject major != 9/10 and TRITON_MLA is the only + # eligible backend @0.24.0). Two measured limits (RTX PRO 6000 + # Blackwell; serving fails identically; re-verify on the next vLLM + # bump): + # 1. FP8 KV + q-heads >= 2 routes to the grouped decode kernel, which + # requests 102400B shared memory vs SM120's 101376B limit -> + # OutOfResources (heads == 1 takes the non-grouped kernel and + # passes; bf16 KV passes). vLLM's overflow guard only drops + # num_stages at BLOCK_DMODEL >= 1024, which the MLA Lk=576 path + # (BLOCK_DMODEL=512) never reaches (triton_decode_attention.py + # :490-532 @0.24.0). Upstream fix in flight: vllm#46728 (open PR, + # num_stages=1 fallback for exactly this tile, fixes vllm#46721). + # 2. Decode batches whose total cached tokens exceed ~2^31/576 raise + # a deterministic illegal memory access (reproduced in isolation + # on a clean GPU): largest passing batch*seq = 2.10M tokens, + # smallest failing = 4.19M, bracketing 2^31 / 576 elements = 3.73M + # — consistent with int32 offset overflow in the Triton decode + # kernel's KV indexing. Affects both KV dtypes (fp8 only at + # heads == 1, where limit 1 does not fire first). Same family on + # SM89 (L40S): 33/400 sampled generation cases, smallest failing + # batch*seq again 4.19M tokens (bf16 KV; fp8-KV combos are not + # declared below SM90). + # FIXME(kernel-limit): on SM120, every DSA case fails: the CUDA sparse + # attention indexer hard-requires DeepGEMM (sparse_attn_indexer.py:468-472 + # @0.24.0), whose fp8 MQA-logits kernels ship for SM90/SM100 only + # (vllm/third_party/deep_gemm .../impls/sm{90,100}_fp8_mqa_logits.cuh) and + # assert "Unsupported architecture" (deepgemm csrc/apis/attention.hpp:184); + # cases with fp8_block linears die even earlier in DeepGEMM's scale-factor + # layout transform ("Unknown SF transformation", layout.hpp:59). Serving + # fails identically. Upstream: vllm#45317 (gap report), TRITON_MLA_SPARSE + # backend for SM8x/11x/12x in vllm#38476/#47629 (open PRs), sparse decode + # in vllm#47527; DeepGEMM SM120 support in DeepGEMM#318 (open PR). + # Re-verify on the next vLLM/DeepGEMM bump. + attn_metadata_dict = {attn_layer_name: attn_metadata} + if indexer_metadata is not None: + attn_metadata_dict[indexer_layer_name] = indexer_metadata + # 0.29: DeepseekV32Attention reads per-layer slot mappings from + # forward_context.slot_mapping (attention.py:303-306@v0.29.0 — a dict + # keyed by layer name); mirror serving's population with the SAME slot + # tensor the metadata was built from. + _slot = getattr(common_attn_metadata, "slot_mapping", None) + _slot_map = dict.fromkeys(attn_metadata_dict, _slot) if _slot is not None else None + point_stack.enter_context(set_forward_context(attn_metadata_dict, vllm_config, slot_mapping=_slot_map)) + try: + with torch.inference_mode(): + _fwd_args = ( + (positions, hidden_states) + if type(attn_module).__name__ == "DeepseekV32Attention" + else (positions, hidden_states, None) + ) + attn_module.forward(*_fwd_args) + except torch.cuda.OutOfMemoryError as e: + print(f" Dry run OOM: {e}") + # Let collect.py record the capacity failure. Returning normally would + # mark a task done even though it emitted no performance row. + raise + except Exception as e: + print(f" Dry run failed: {e}") + traceback.print_exc() + # Propagate to collect.py's worker so the failure is recorded in the + # error queue. Swallowing here lets a whole op complete with zero + # rows while the pipeline reports success. + raise + + _lap(point_phase, "dry_run") # 5. Benchmark def kernel_func(): - attn_module.forward(positions, hidden_states, None) - - # DSA context previously bypassed graphs to avoid retained scratch pools. - # Keep graph timing for small as well as large cases: eager launch gaps can - # dominate small module latency. The helper resets each graph, and the - # ExitStack releases vLLM's workspace after its forward contexts unwind. - with benchmark_with_power( - device=torch.device(device), - kernel_func=kernel_func, - num_warmups=warming_up, - num_runs=test_ite, - repeat_n=1, - use_cuda_graph=True, - allow_graph_fail=False, - ) as results: - pass - - if results.get("used_cuda_graph") is not True: - raise RuntimeError("vLLM MLA/DSA module requires CUDA graph timing; refusing to publish eager timing") - latency = results["latency_ms"] - - # 6. Log results — schema aligned with TRT-LLM - if is_context: - isl = seq_len - step = prefix_len - else: - isl = 1 - # BatchSpec.seq_lens includes this forward's one new token. - # Both module loaders reconstruct total KV length as isl + step. - step = seq_len - 1 - - op_name = f"{attn_type}_{phase}_module" + ("_skip_indexer" if skip_indexer else "") - - # Record architecture to distinguish different DSA models in the perf CSV. - # perf_database uses this as a dict key when loading data. - # Aligns with sdk/models.py which uses architectures[0] throughout. - hf_cfg = vllm_config.model_config.hf_config - architecture = getattr(hf_cfg, "architectures", [getattr(hf_cfg, "model_type", "unknown")])[0] - mla_layer = attn_module.mla_attn.mla_attn - backend_name = _mla_backend_name(mla_layer, attn_type, is_context, attn_metadata) - actual_kv_cache_dtype = "fp8" if mla_layer.kv_cache_dtype.startswith("fp8") else "bfloat16" - - log_perf( - item_list=[ - { - "model": model_path, - "architecture": architecture, - "mla_dtype": "bfloat16" if compute_dtype == "bfloat16" else compute_dtype, - "kv_cache_dtype": actual_kv_cache_dtype, - "gemm_type": "bfloat16" if gemm_type == "bfloat16" else gemm_type, - "num_heads": num_heads, - "batch_size": batch_size, - "isl": isl, - "tp_size": 1, - "step": step, - "latency": f"{latency:.4f}", - } - ], - framework="VLLM", - version=vllm_version, - device_name=torch.cuda.get_device_name(device), - op_name=op_name, - kernel_source=backend_name, - perf_filename=perf_filename, - power_stats=results["power_stats"], - ) + attn_module.forward(*_fwd_args) + + def measure_and_log(skip_indexer): + if has_reuse: + skip_owner.skip_topk = skip_indexer + + # Graph timing for every MLA/DSA case. DSA context used to run eagerly to dodge a + # private-pool retention (~18 GiB flashmla-sparse scratch per capture, saturating + # at ~146 GiB per worker); that bought small cases a ~0.6 ms launch floor — 0.30 + # rows at isl <= 1024 measured 1.6-8x the graph-timed values of the same cells + # (b200/l40s vs #404's B200 tables, 2026-10-10). benchmark_with_power + # (explicit_graph_cleanup) resets each graph and the ExitStack releases vLLM's + # workspace after the forward contexts unwind. FIXME(re-verify): confirm on 0.30 + # with a small->large DSA context sweep that the retention is gone before a campaign. + with benchmark_with_power( + device=torch.device(device), + kernel_func=kernel_func, + num_warmups=warming_up, + # The point's dry run already executed this forward; three eager + # passes cover lazily initialised paths (eager_warmup was ~40% of a + # measurement's wall time at ten). + num_eager_warmups=3, + num_runs=test_ite, + repeat_n=1, + use_cuda_graph=True, + allow_graph_fail=False, + ) as results: + pass + _lap(point_phase, "measure_skip" if skip_indexer else "measure_full") + # Protocol phases of this measurement (eager_warmup, capture, graph_warmup, + # timed, cleanup) on the point's timing line, prefixed by the measurement. + _bench_prefix = "skip" if skip_indexer else "full" + for _phase_name, _phase_secs in (results.get("phase_seconds") or {}).items(): + point_phase[f"{_bench_prefix}_{_phase_name}"] = _phase_secs + + if results.get("used_cuda_graph") is not True: + raise RuntimeError("vLLM MLA/DSA module requires CUDA graph timing; refusing to publish eager timing") + latency = results["latency_ms"] + + # 6. Log results — schema aligned with TRT-LLM + if is_context: + isl = seq_len + step = prefix_len + else: + isl = 1 + # BatchSpec.seq_lens includes this forward's one new token. + # Both module loaders reconstruct total KV length as isl + step. + step = seq_len - 1 + + op_name = f"{attn_type}_{phase}_module" + ("_skip_indexer" if skip_indexer else "") + backend_name = _mla_backend_name(mla_layer, attn_type, is_context, attn_metadata) + + log_perf( + item_list=[ + { + "model": model_path, + "architecture": architecture, + "mla_dtype": "bfloat16" if compute_dtype == "bfloat16" else compute_dtype, + "kv_cache_dtype": actual_kv_cache_dtype, + "gemm_type": "bfloat16" if gemm_type == "bfloat16" else gemm_type, + "num_heads": num_heads, + "batch_size": batch_size, + "isl": isl, + "tp_size": 1, + "step": step, + "latency": f"{latency:.4f}", + } + ], + framework="VLLM", + version=vllm_version, + device_name=torch.cuda.get_device_name(device), + op_name=op_name, + kernel_source=backend_name, + perf_filename=perf_filename, + power_stats=results["power_stats"], + ) - print( - f" [{phase}] b={batch_size}, s={seq_len}, heads={num_heads}, " - f"prefix={prefix_len}, gemm={gemm_type}, compute={compute_dtype}, " - f"kv={kv_cache_dtype}, backend={backend_name}: {latency:.4f} ms" - ) + print( + f" [{phase}] b={batch_size}, s={seq_len}, heads={num_heads}, " + f"prefix={prefix_len}, gemm={gemm_type}, compute={compute_dtype}, " + f"kv={kv_cache_dtype}, backend={backend_name}: {latency:.4f} ms" + ) - return latency + return latency - full_latency = measure_and_log(False) - if has_reuse: - measure_and_log(True) - return full_latency + full_latency = measure_and_log(False) + if has_reuse: + measure_and_log(True) + return full_latency + + latencies: list[float] = [] + failures: list[tuple[tuple[int, int], BaseException]] = [] + for index, (point_seq_len, point_prefix) in enumerate(points): + _t_point = _t_last = time.perf_counter() + point_phase: dict[str, float] = dict(_setup_phase) if index == 0 else {} + try: + with ExitStack() as point_stack: + latencies.append(_run_point(point_stack, point_seq_len, point_prefix, point_phase)) + except Exception as exc: + if _is_fatal_cuda_error(exc) or len(points) == 1: + raise + # The traceback holds _run_point's frame (KV cache, indexer cache, metadata, + # activations): clear it, or the largest-first series keeps that memory + # pinned and the smaller points that would fit fail too. + traceback.clear_frames(exc.__traceback__) + failures.append(((point_seq_len, point_prefix), exc)) + print( + f" [{variant} module] {phase} b={batch_size}, s={point_seq_len}, prefix={point_prefix}: " + f"point failed ({type(exc).__name__}: {str(exc)[:200]}); the series continues" + ) + gc.collect() + torch.cuda.empty_cache() + continue + point_phase["indexer_fill"] = globals().get("_PHASE_TIMINGS", {}).get("indexer_fill", 0.0) + print( + f" [{variant} module timing] {phase} b={batch_size}, s={point_seq_len}, prefix={point_prefix}, " + f"heads={num_heads}, gemm={gemm_type}, kv={kv_cache_dtype}: " + + " ".join(f"{name}={seconds:.2f}s" for name, seconds in point_phase.items()) + + f" total={time.perf_counter() - _t_point:.2f}s" + ) + if len(points) > 1: + print( + f" [{variant} module series] {phase} b={batch_size}, prefix={prefix_desc}, heads={num_heads}, " + f"gemm={gemm_type}, kv={kv_cache_dtype}: {len(latencies)}/{len(points)} points in " + f"{time.perf_counter() - _t_series:.1f}s" + ) + if failures: + if len(failures) == 1: + raise failures[0][1] + summary = "; ".join(f"s={s} prefix={p}: {type(exc).__name__}: {str(exc)[:200]}" for (s, p), exc in failures) + raise RuntimeError(f"{len(failures)}/{len(points)} point(s) of the module series failed: {summary}") + return latencies def run_mla_module_worker( - seq_len: int, + seq_len: int | list[int], batch_size: int, num_heads: int, kv_cache_dtype: str, @@ -1226,12 +1858,37 @@ def run_mla_module_worker( gemm_type: str, model_path: str, attn_type: str, - prefix_len: int = 0, + prefix_len: int | list[int] = 0, *, perf_filename: str, device: str = "cuda:0", ): - """Worker-compatible positional wrapper used by collector/collect.py.""" + """Worker-compatible positional wrapper used by collector/collect.py. + + ``seq_len`` is one point (the CLI, single-point fills) or the case's list of + points measured on one module (the registry plan). + """ + # FIXME(kernel-limit): on SM100/103 FlashInfer's trtllm-gen sparse-MLA kernels start at tileSizeQ 8 + # (= per-rank heads at q_len 1): DSA heads 1/2/4 (TP 32/64/128 of a 128-head model) raise "Missing + # TRTLLM-GEN kernel" (csrc/trtllm_fmha_kernel_launcher.cu:318 @flashinfer 0.6.18, vLLM 0.30.0's pin), + # 68,435 b200 cases after a full module load each (job 471043976). Serving fails identically. + if attn_type == "dsa" and get_sm_version() in (100, 103) and num_heads < 8: + raise RuntimeError( + "FIXME(kernel-limit): FlashInfer trtllm-gen sparse-MLA on SM100/103 has no kernel for per-rank " + f"heads < 8 (num_heads={num_heads}; trtllm_fmha_kernel_launcher.cu:318 @flashinfer 0.6.18)" + ) + # FIXME(kernel-limit): on SM100/103 with FP8 KV, vLLM selects FLASHINFER_MLA_SPARSE for DSA and + # FlashInfer's auto backend has nothing for 64 < per-rank heads < 128: trtllm-gen refuses the head + # range, cute-dsl refuses sparse_mla_top_k ("auto: no backend supports this configuration", 11,442 + # cases = every heads-96 fp8 cell, both phases, job 471043976). Serving fails identically (the same + # gap closes Kimi-K3 tp1 on sglang: finding sm100_probe_failures_recheck_2026_10_04 (i)). Raise + # before the module is built so the cell is a classified failure, not a 20 s load + crash. + if attn_type == "dsa" and "fp8" in str(kv_cache_dtype) and get_sm_version() in (100, 103) and 64 < num_heads < 128: + raise RuntimeError( + "FIXME(kernel-limit): vLLM FLASHINFER_MLA_SPARSE on SM100/103 has no decode backend for " + f"64 < per-rank heads < 128 with FP8 KV (num_heads={num_heads}); trtllm-gen rejects the head " + "range and cute-dsl rejects sparse_mla_top_k (flashinfer mla/_core.py:4271); serving fails identically" + ) # Serving executes model forward under inference mode (vLLM dd10e03f9, # v1/worker/gpu_model_runner.py:4069-4070). Keep module initialization, # dry runs, graph warmup and timing in the SAME mode: FlashInfer 0.6.13 @@ -1321,7 +1978,7 @@ def main(): parser.add_argument( "--gemm-type", type=str, - choices=["bfloat16", "fp8", "fp8_block", "nvfp4"], + choices=["bfloat16", "fp8_block", "nvfp4"], default=None, help="GEMM quantisation type for linear layers (default: run all supported by GPU)", ) diff --git a/python/aisimulate/collector/vllm/collect_mla_module_027.py b/python/aisimulate/collector/vllm/collect_mla_module_027.py deleted file mode 100644 index f7b7919f50..0000000000 --- a/python/aisimulate/collector/vllm/collect_mla_module_027.py +++ /dev/null @@ -1,36 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""vLLM 0.27.0 lane for the MLA module collector. - -`collector.vllm.collect_mla_module` stays pinned to the manifest default -(0.24.0) because it also hosts the DSA ops whose family (`sparse_attention`) -collects on that runtime. The `mla` family moved to the 0.27.0 K3 serving -lane image; per the one-runtime-per-module invariant those ops resolve here. - -The bodies are thin re-dispatch wrappers around the base collector: the -0.27.0 audit (GB300 probe, 2026-08-17) confirmed -`DeepseekV2MLAAttention.__init__` params, the `glm_moe_dsa` config-registry -gap patch, `backend_supports_prefill_query_quantization`, and the prefill -selector surface are all unchanged from the 0.24.0 citations documented in -collect_mla_module.py, so the wrappers change neither WHICH kernel runs nor -HOW a case is timed. -""" - -# The historical filename also serves the source/runtime-qualified 0.25.0 -# MLA family override; each campaign manifest still selects one exact runtime. -__compat__ = "vllm>=0.25.0,<=0.27.0,!=0.25.1,!=0.26.0" - -from collector.vllm import collect_mla_module as _impl - - -def get_mla_context_module_test_cases(): - return _impl.get_mla_context_module_test_cases() - - -def get_mla_generation_module_test_cases(): - return _impl.get_mla_generation_module_test_cases() - - -def run_mla_module_worker(*args, **kwargs): - return _impl.run_mla_module_worker(*args, **kwargs) diff --git a/python/aisimulate/collector/vllm/collect_moe.py b/python/aisimulate/collector/vllm/collect_moe.py index 1e1483ff46..9acdefca37 100644 --- a/python/aisimulate/collector/vllm/collect_moe.py +++ b/python/aisimulate/collector/vllm/collect_moe.py @@ -72,9 +72,11 @@ # 1968047). Existing framework construction/dispatch is unchanged. Invalid # model/quant shapes continue to raise and remain recorded failures; this is # API compatibility, not a claim that every Cartesian shape is supported. -# 0.30.0 source audit (2026-09-30; tag v0.30.0 == ced6857afa0e vs v0.27.1), -# the baseline of the 0.31.0 audit below. 0.28.0, 0.29.0 and 0.30.0 itself -# are not admitted: no consumer reads their tables. +# 0.30.0 source audit (2026-09-30; tag v0.30.0 == ced6857afa0e vs v0.27.1): +# the collector default (framework_manifest.yaml; the 2026-10 campaign +# tables) and the baseline of the 0.31.0 audit below, which admits the +# GLM-5.3-Flash model-pinned run. 0.28.0/0.29.0 stay admitted from the +# 0.27.1 -> 0.30.0 upgrade path. # FusedMoEFactory (layer.py:88-131 @0.30.0, :99-140 @0.27.1): every kwarg # passed here is unchanged; 0.30.0 only adds defaulted params. swiglu_limit # exists at both (:108 / :119), so other models keep passing None. @@ -102,10 +104,14 @@ # TrtLlmNvFp4ExpertsMonolithic (trtllm_nvfp4_moe.py:486-508, clamp # :90-105,:175-181). Serving logs of both GLM checkpoints on this runtime # report "Using FLASHINFER_TRTLLM" for Fp8 and NvFp4 MoE. -# Known non-parity (unchanged from 0.27.1): no FlashInfer autotune pass -# before timing; the gate GEMM and shared expert are separate ops. +# Known non-parity: the gate GEMM and shared expert are separate ops. +# (FlashInfer lanes are timed on the autotuned tactic like serving; see +# helper.flashinfer_autotuned.) # 0.31.0 audit (2026-10-07, source-only: tag v0.31.0 == db9527a4 vs v0.30.0 -# ced6857a; no releases in between). GLM-5.3-Flash now runs STOCK 0.31.0 +# ced6857a; no releases in between). The moe family runs on 0.31.0 +# (framework_manifest.yaml family pin, 2026-10-30), so this file targets +# exactly 0.31.0. One release per file (owner policy 2026-09-20); earlier code is `git log -- ` away. +# GLM-5.3-Flash now runs STOCK 0.31.0 # (no overlay). Construction surface unchanged: FusedMoEFactory # (layer.py:88-141) keeps every kwarg passed here (swiglu_limit :108, # activation_situ_beta :111, router_logits_dtype :125, @@ -160,7 +166,7 @@ # K3 SiTU: FusedMoEFactory exposes activation_situ_beta on stock 0.30.0 and # 0.31.0, so the probe below passes the betas natively there; the # situ-as-silu approximation applies only to older stock builds. -__compat__ = "vllm>=0.24.0,<=0.31.0,!=0.25.1,!=0.26.0,!=0.27.0,!=0.28.0,!=0.29.0,!=0.30.0" +__compat__ = "vllm==0.31.0" import contextlib import json @@ -178,6 +184,7 @@ from collector.helper import ( balanced_logits, benchmark_with_power, + flashinfer_autotuned, get_sm_version, log_perf, power_law_logits_v3, @@ -632,7 +639,25 @@ def gemma4_routing_function(hidden_states, gating_output, topk, renormalize): # the Mxfp4Config path below. quant_config = CompressedTensorsConfig.from_config(checkpoint_qc) else: + # OCP "mxfp4" checkpoints of model_type gpt_oss are OVERRIDDEN by + # vLLM to the gpt_oss_mxfp4 method (GptOssMxfp4Config. + # override_quantization_method, quantization/mxfp4.py:119-136 + # @0.29.0). Its GptOssMxfp4MoEMethod selects through the gpt-oss + # priority list (fused_moe/oracle/mxfp4.py:584-602 @0.29.0) — TRITON + # matmul_ogs on SM90, as the gpt-oss-120b serving probe shows — + # while the generic Mxfp4Config resolves to Mxfp4MoEMethod -> MARLIN + # (path_diff moe_mxfp4 vs gpt-oss-120b was diverged on exactly this, + # 2026-09-23). Mirror the framework's own override here + # (GptOssMxfp4Config, quantization/mxfp4.py:110 @0.31.0). + model_cfg = _load_model_moe_config(model_name) quant_config = Mxfp4Config() + if ( + checkpoint_qc.get("quant_method") in ("mxfp4", "gpt_oss_mxfp4") + and model_cfg.get("model_type") == "gpt_oss" + ): + from vllm.model_executor.layers.quantization.mxfp4 import GptOssMxfp4Config + + quant_config = GptOssMxfp4Config() elif moe_type == "w4a8_mxfp4_mxfp8": # Native DeepSeek-V4 (expert_dtype=fp4) serving path: vLLM overrides # the checkpoint's fp8 quant_method to DeepseekV4FP8Config @@ -955,8 +980,11 @@ def run_single_iteration(): forward_context.moe_layer_index = 0 moe_module(hidden_states, router_logits) + # FlashInfer-backed experts (trtllm-gen nvfp4/mxfp4, cutlass, cutedsl) are + # timed on the autotuned tactic like serving; see helper.flashinfer_autotuned. with ( set_forward_context({}, vllm_config), + flashinfer_autotuned(num_tokens, run_single_iteration, label=source), benchmark_with_power( device=device, kernel_func=run_single_iteration, diff --git a/python/aisimulate/collector/vllm/collect_moe_xpu.py b/python/aisimulate/collector/vllm/collect_moe_xpu.py index 59a859162d..7806c81c1d 100644 --- a/python/aisimulate/collector/vllm/collect_moe_xpu.py +++ b/python/aisimulate/collector/vllm/collect_moe_xpu.py @@ -293,7 +293,10 @@ def run_single_iteration(): ) def run_iterations(): - # Use benchmark_with_power context manager + # Use benchmark_with_power context manager. XPU graph capture follows + # vLLM's own VLLM_XPU_ENABLE_XPU_GRAPH (xpu_graph_measure_enabled), and + # a failed capture fails the case (allow_graph_fail stays at its False + # default) instead of silently timing eager. with benchmark_with_power( device=device, kernel_func=run_single_iteration, @@ -301,7 +304,6 @@ def run_iterations(): num_runs=num_runs, repeat_n=1, use_cuda_graph=xpu_graph_measure_enabled(), - allow_graph_fail=False, # graph mandatory; capture failure fails the case ) as results: pass diff --git a/python/aisimulate/collector/vllm/collect_msa_module.py b/python/aisimulate/collector/vllm/collect_msa_module.py index 7dff409487..713d08bb4e 100644 --- a/python/aisimulate/collector/vllm/collect_msa_module.py +++ b/python/aisimulate/collector/vllm/collect_msa_module.py @@ -3,14 +3,11 @@ # Portions adapted from vLLM dd10e03f95f94edbea1975c67ace3a35ec9a8a40. # Copyright contributors to the vLLM project. Modified for collector execution. -# MiniMax-M3 support landed in vLLM 0.24.0 (vllm/models/minimax_m3/ + -# MinimaxM3QKVParallelLinearWithIndexer + fused_minimax_m3_qknorm_rope_kv_insert); -# this collector follows the 0.24.0 APIs and is pinned to the manifest -# default runtime (vllm collectors pin exactly; see -# test_active_cuda_vllm_collectors_are_exactly_pinned_to_manifest_version). -# B200 0.25.0 validation: job 1968755 passed 8 context and 6 generation -# cases after matching serving inference mode, positions, and THK top-k layout. -__compat__ = "vllm>=0.24.0,<=0.25.0" +# Versions are upgraded IN PLACE (owner policy 2026-09-20): this file always +# targets the manifest pin; prior-version code is `git log -- ` +# away and prior-version DATA stays permanent under its version key. The +# 0.24->0.29 adaptations each carry a serving citation @0.29.0. +__compat__ = "vllm==0.30.0" """ MSA Module Collector for vLLM — MiniMax-M3 sparse-attention benchmarking. @@ -80,6 +77,7 @@ get_mla_module_model_specs, get_mla_module_precision_specs, get_mla_module_sweep_spec, + get_msa_head_counts, ) from collector.helper import _resolve_local_model_path, benchmark_with_power, get_sm_version, log_perf from collector.registry_types import PerfFile @@ -166,8 +164,9 @@ def get_context_test_cases(): """ cases = [] sweep = get_mla_module_sweep_spec("vllm") + head_counts = get_msa_head_counts("vllm") # native // tp of the declared shards, not the shared head grid for compute_dtype, kv_dtype, gemm_type in _get_precision_combos("context"): - for num_heads in sweep.inner_sweep_head_counts: + for num_heads in head_counts: for b in sweep.context_batch_sizes: for s in sweep.context_sequence_lengths: if b * s > sweep.context_max_tokens: @@ -199,8 +198,9 @@ def get_generation_test_cases(): budget = None if total_memory is None else int(total_memory * _MEMORY_BUDGET_SAFETY_FACTOR) considered = 0 dropped = 0 + head_counts = get_msa_head_counts("vllm") for compute_dtype, kv_dtype, gemm_type in _get_precision_combos("generation"): - for num_heads in sweep.inner_sweep_head_counts: + for num_heads in head_counts: for b in sweep.generation_batch_sizes: for s in sweep.generation_sequence_lengths: if b * s > sweep.generation_max_tokens: @@ -302,6 +302,30 @@ def _create_gemm_quant_config(gemm_type: str): _INDEX_CACHE_LAYER_NAME = "model.layers.0.self_attn.attn.index_cache" +def msa_decode_scorer_kernel_limit( + num_idx_heads: int, num_heads: int, *, is_context: bool, query_len: int | None +) -> str | None: + """Classified kernel-limit message when the M3 DECODE indexer scorer + cannot tile this index-head count, else None. + + The scorer tiles BLOCK_SIZE_HQ = num_idx_heads * BLOCK_SIZE_Q through + tl.arange (power of 2 required). It runs for every generation cell and + for a one-token context query only; prefill queries take the other + kernel, which the 0.29 full-grid sweep collected at 48/96 heads + (904/976 rows per head, the 72 misses being exactly the isl=1 cells). + """ + decode_scorer_runs = (not is_context) or (query_len is not None and query_len <= 1) + if decode_scorer_runs and num_idx_heads & (num_idx_heads - 1): + return ( + f"FIXME(kernel-limit): num_idx_heads={num_idx_heads} (from num_heads={num_heads}) " + "is not a power of 2; MiniMax-M3 sparse decode indexer Triton kernel requires " + "power-of-2 tl.arange tiles (models/minimax_m3/common/ops/" + "sparse_attn.py@0.29.0) and serving tp shards can never produce " + "this count" + ) + return None + + def _create_msa_attention_module( model_path: str, num_heads: int, @@ -311,10 +335,15 @@ def _create_msa_attention_module( max_batch_size: int, is_context: bool, device: str = "cuda:0", + query_len: int | None = None, ): """Create a ``MiniMaxM3SparseAttention`` module from vLLM's own modeling code. + ``query_len`` is the per-request query length of the case (context + ``seq_len``; ``1`` for generation) — it decides whether the decode-path + kernel-limit guard below applies. + Loads the real HF config from model_path, applies the TP-shard head emulation in-memory, and constructs the exact module serving builds for sparse layers (``MiniMaxM3DecoderLayer.__init__``, @@ -410,6 +439,28 @@ def _create_msa_attention_module( # (nvidia/model.py:769-781). sparse_cfg = dict(hf_config.sparse_attention_config) sparse_cfg["sparse_num_index_heads"] = num_kv_heads + + # KERNEL LIMIT GUARD (verified on the 0.29 full-grid sweep, 2026-09-21): + # the M3 DECODE indexer-score Triton kernel tiles BLOCK_SIZE_HQ = + # num_idx_heads * BLOCK_SIZE_Q through tl.arange, whose range must be a + # power of 2 (models/minimax_m3/common/ops/sparse_attn.py — the sweep's + # heads 48/96 shard to idx_heads 3/6 and fail compilation + # deterministically). Serving cannot reach these points either: real tp + # in {1,2,4,8} on the native 4 index heads yields idx_heads in {4,2,1}, + # always a power of 2. Raise the classified error up front instead of + # burning a Triton compile per case. + # SCOPE (0.30 full-grid sweep, 2026-09-28): the limit is the decode + # path's. On 0.29 the 48/96-head failures were every generation cell and + # ONLY the isl=1 context cells (the single-query-token context case runs + # the decode scorer): 736/736 generation, 72/976 context per head, while + # context isl >= 16 collected 904/904 rows per head. A guard on the head + # count alone dropped those 1808 prefill rows on 0.30 — so it now fires + # only where the decode scorer runs (generation, or a one-token query). + _limit = msa_decode_scorer_kernel_limit( + int(sparse_cfg["sparse_num_index_heads"]), num_heads, is_context=is_context, query_len=query_len + ) + if _limit: + raise ValueError(_limit) hf_config.sparse_attention_config = sparse_cfg # Reserved top-k indices buffer shared by the indexer and the attend — @@ -500,7 +551,9 @@ def _rebase_block_table_and_slots(common_attn_metadata, block_size: int): block_table += 1 query_start_loc = common_attn_metadata.query_start_loc_cpu - context_lens = common_attn_metadata.num_computed_tokens_cpu + # 0.30: property removed; seq_lens - query_lens on device + # (utils.py CommonAttentionMetadata.compute_num_computed_tokens@v0.30.0) + context_lens = common_attn_metadata.compute_num_computed_tokens().cpu() slot_mapping = common_attn_metadata.slot_mapping device = slot_mapping.device for i in range(common_attn_metadata.num_reqs): @@ -588,38 +641,44 @@ def _create_kv_caches_and_metadata( ) num_blocks = _rebase_block_table_and_slots(common_attn_metadata, block_size) - # Main paged K/V cache: shape from the layer's own backend - # ((num_blocks, 2, block_size, num_kv_heads, head_size) — MiniMaxM3 - # SparseBackend.get_kv_cache_shape, common/sparse_attention.py:90-98 - # @v0.24.0; NHD stride order is the natural layout, :100-114), dtype - # from the layer's KV-cache spec (nvidia/model.py:540-549@v0.24.0: - # spec.dtype = kv_cache_dtype_str_to_dtype(cache_dtype) — bf16 for - # "auto", uint8 storage for "fp8", utils/torch_utils.py:38,394-400). - # Serving allocates the same way: a raw per-layer buffer viewed with the - # spec dtype and backend shape (gpu_model_runner - # ._reshape_kv_cache_tensors:7081-7180@v0.24.0); the impl then - # reinterprets the uint8 storage as the platform fp8 (e4m3) at use - # (common/sparse_attention.py:298-306, view at :352). + # Main paged K/V cache. 0.29 removed AttentionBackend.get_kv_cache_shape; + # the shape contract moved into the generic spec pipeline: serving + # allocates a flat byte buffer and views each layer as the 4D logical + # [B, H, N, C] page, C in bytes (compute_layer_kv_cache_shape_bytes, + # v1/kv_cache_interface.py:253-269@v0.29.0; the per-layer element view is + # cache_logical.view(spec.dtype), create_kv_cache_views :371-376). For + # M3's main cache the spec is FullAttentionSpec (nvidia/model.py:556-565) + # so the element view is [num_blocks, num_kv_heads, block_size, + # 2*head_size] — merged K|V on the last dim: the impl splits it back with + # k_cache, v_cache = kv_cache.split(head_size, dim=-1) + # (nvidia/sparse_attention_msa.py:277) and the fused insert kernel reads + # the block size off dim 2 (kv_cache.size(2), nvidia/model.py:635). + # Dtype semantics are unchanged: spec.dtype is bf16 for "auto", uint8 + # storage for "fp8"; the impl reinterprets uint8 as e4m3 at use + # (nvidia/sparse_attention_msa.py:221-222). + from vllm.v1.kv_cache_interface import compute_layer_kv_cache_shape_bytes + backend_cls = attn_module.get_attn_backend() kv_cache_spec = attn_module.get_kv_cache_spec(vllm_config) kv_cache = torch.zeros( - backend_cls.get_kv_cache_shape(num_blocks, block_size, kv_cache_spec.num_kv_heads, kv_cache_spec.head_size), - dtype=kv_cache_spec.dtype, + compute_layer_kv_cache_shape_bytes(kv_cache_spec, num_blocks), + dtype=torch.uint8, device=torch_device, - ) + ).view(kv_cache_spec.dtype) - # Indexer key side-cache: (num_blocks, block_size, head_size), dtype from - # the index cache module (bf16 for the default indexer_kv_dtype "bf16", - # config/attention.py:55; MiniMaxM3IndexerBackend.get_kv_cache_shape, - # common/indexer.py:91-99@v0.24.0). + # Indexer key side-cache: MLAAttentionSpec, key-only — one head slot, one + # vector per token (common/indexer.py:144-151@v0.29.0), logical + # [num_blocks, 1, block_size, head_size] after the dtype view. Serving + # hands this 4D view to the layer's bind_kv_cache, which squeezes the + # H=1 dim (common/indexer.py:140-142); binding below mirrors that. index_cache_layer = attn_module.indexer.index_cache index_backend_cls = index_cache_layer.get_attn_backend() index_spec = index_cache_layer.get_kv_cache_spec(vllm_config) index_kv_cache = torch.zeros( - index_backend_cls.get_kv_cache_shape(num_blocks, block_size, 1, index_spec.head_size), - dtype=index_cache_layer.dtype, + compute_layer_kv_cache_shape_bytes(index_spec, num_blocks), + dtype=torch.uint8, device=torch_device, - ) + ).view(index_spec.dtype) main_builder = backend_cls.get_builder_cls()(kv_cache_spec, [_ATTN_LAYER_NAME], vllm_config, torch_device) attn_metadata = main_builder.build( @@ -692,6 +751,7 @@ def run_msa_module( max_batch_size=batch_size, is_context=is_context, device=device, + query_len=seq_len if is_context else 1, ) # 2. Create KV caches + metadata via the framework's builders. @@ -706,14 +766,16 @@ def run_msa_module( device=device, ) - # 2b. Bind the caches to the registered layers — the benchmark analog of - # serving's bind_kv_cache, which assigns each allocated per-layer tensor - # to forward_context[layer_name].kv_cache (v1/worker/utils.py:462-530 - # @v0.24.0; the model reads self.kv_cache / index_cache.kv_cache directly, - # nvidia/model.py:606-607). + # 2b. Bind the caches to the registered layers exactly as serving does: + # forward_context[layer_name].bind_kv_cache(kv_cache) + # (v1/worker/utils.py:622@v0.29.0). The main layer's bind is the + # AttentionLayerBase default (stores the view as-is, + # model_executor/layers/attention_layer_base.py:26-32); the indexer's + # override squeezes the H=1 head dim, [B,1,N,C] -> [B,N,C] + # (common/indexer.py:140-142). forward_ctx = vllm_config.compilation_config.static_forward_context - forward_ctx[_ATTN_LAYER_NAME].kv_cache = kv_cache - forward_ctx[_INDEX_CACHE_LAYER_NAME].kv_cache = index_kv_cache + forward_ctx[_ATTN_LAYER_NAME].bind_kv_cache(kv_cache) + forward_ctx[_INDEX_CACHE_LAYER_NAME].bind_kv_cache(index_kv_cache) # 2c. Prove the requested KV precision actually reached the framework # path — a benchmark that silently ran bf16 under an "fp8" label would be diff --git a/python/aisimulate/collector/vllm/registry.py b/python/aisimulate/collector/vllm/registry.py index 3e21c25c08..3967de160c 100644 --- a/python/aisimulate/collector/vllm/registry.py +++ b/python/aisimulate/collector/vllm/registry.py @@ -25,6 +25,9 @@ get_func="get_computescale_test_cases", run_func="run_computescale", perf_filename=PerfFile.COMPUTESCALE, + # one measurement, two tables: the static quant row goes to + # scale_matrix_perf; finalize binds both to this producer + extra_perf_filenames=(PerfFile.SCALE_MATRIX,), ), OpEntry( op="attention_context", @@ -70,14 +73,14 @@ ), OpEntry( op="mla_context_module", - module="collector.vllm.collect_mla_module_027", + module="collector.vllm.collect_mla_module", get_func="get_mla_context_module_test_cases", run_func="run_mla_module_worker", perf_filename=PerfFile.MLA_CONTEXT_MODULE, ), OpEntry( op="mla_generation_module", - module="collector.vllm.collect_mla_module_027", + module="collector.vllm.collect_mla_module", get_func="get_mla_generation_module_test_cases", run_func="run_mla_module_worker", perf_filename=PerfFile.MLA_GENERATION_MODULE, @@ -155,13 +158,6 @@ run_func="run_dsv4_sparse_kernel_worker", perf_filename=PerfFile.DSV4_PAGED_MQA_LOGITS_MODULE, ), - OpEntry( - op="dsv4_hca_attn_module", - module="collector.vllm.collect_dsv4_attn", - get_func="get_dsv4_hca_attn_test_cases", - run_func="run_dsv4_sparse_kernel_worker", - perf_filename=PerfFile.DSV4_HCA_ATTN_MODULE, - ), OpEntry( op="mhc_module", module="collector.vllm.collect_mhc_module", diff --git a/python/aisimulate/collector/vllm/utils.py b/python/aisimulate/collector/vllm/utils.py index 4e4885c1ad..b19ae0d39d 100644 --- a/python/aisimulate/collector/vllm/utils.py +++ b/python/aisimulate/collector/vllm/utils.py @@ -116,13 +116,17 @@ def create_common_attn_metadata( # Calculate max query length max_query_len = max(batch_spec.query_lens) + # Serving population site: vllm/v1/worker/gpu/attn_utils.py:456-475@v0.30.0 + # passes seq_lens + seq_lens_cpu_upper_bound (no private *_cpu fields — the + # 0.29 `_seq_lens_cpu` / `_num_computed_tokens_cpu` were removed in 0.30; + # computed-token counts derive from seq_lens and query_start_loc inside the + # dataclass). Version bump 2026-09-27 (opharness upgrade_op vllm 0.30.0). + del num_computed_tokens_cpu return CommonAttentionMetadata( query_start_loc=query_start_loc, query_start_loc_cpu=query_start_loc_cpu, seq_lens=seq_lens, seq_lens_cpu_upper_bound=seq_lens_cpu, - _seq_lens_cpu=seq_lens_cpu, - _num_computed_tokens_cpu=num_computed_tokens_cpu, num_reqs=batch_spec.batch_size, num_actual_tokens=num_tokens, max_query_len=max_query_len, @@ -164,7 +168,7 @@ def create_vllm_config( max_model_len: int = 1024, dtype: Union[ModelDType, torch.dtype] = "auto", num_gpu_blocks: int = 1000, - block_size: int = 16, + block_size: int | None = 16, max_num_seqs: int = 256, max_num_batched_tokens: int = 8192, enable_chunked_prefill: bool = True, @@ -195,6 +199,11 @@ def create_vllm_config( hf_overrides=hf_overrides or {}, ) + # block_size=None is NOT "16": pydantic's CacheConfig marks any explicit value as + # user_specified_block_size, and the selector (v1/attention/selector.py:131-135) then + # filters backends by it exactly as --block-size would. Serving without --block-size + # selects unconstrained and derives the page from the chosen backend afterwards + # (utils.framework_kv_block_size); pass None to get that behaviour. cache_config = CacheConfig( block_size=block_size, cache_dtype="fp8" if use_fp8_kv_cache else "auto", @@ -340,9 +349,13 @@ def create_and_prepopulate_kv_cache_mla( MLA KV cache tensor """ batch_size = len(kv_c_contexts) - seq_lens = common_attn_metadata.seq_lens_cpu + # 0.30 removed the seq_lens_cpu property; the CPU view serving carries is + # seq_lens_cpu_upper_bound (vllm/v1/worker/gpu/attn_utils.py:460@v0.30.0) + seq_lens = common_attn_metadata.seq_lens_cpu_upper_bound query_lens = common_attn_metadata.query_start_loc_cpu[1:] - common_attn_metadata.query_start_loc_cpu[:-1] - context_lens = common_attn_metadata.num_computed_tokens_cpu + # 0.30: property removed; seq_lens - query_lens on device + # (utils.py CommonAttentionMetadata.compute_num_computed_tokens@v0.30.0) + context_lens = common_attn_metadata.compute_num_computed_tokens().cpu() block_table = common_attn_metadata.block_table_tensor slot_mapping = common_attn_metadata.slot_mapping @@ -452,9 +465,13 @@ def create_kv_cache_and_block_mappings( Tuple of the empty KV cache and flattened history slot mapping """ batch_size = common_attn_metadata.num_reqs - seq_lens = common_attn_metadata.seq_lens_cpu + # 0.30 removed the seq_lens_cpu property; the CPU view serving carries is + # seq_lens_cpu_upper_bound (vllm/v1/worker/gpu/attn_utils.py:460@v0.30.0) + seq_lens = common_attn_metadata.seq_lens_cpu_upper_bound query_lens = common_attn_metadata.query_start_loc_cpu[1:] - common_attn_metadata.query_start_loc_cpu[:-1] - context_lens = common_attn_metadata.num_computed_tokens_cpu + # 0.30: property removed; seq_lens - query_lens on device + # (utils.py CommonAttentionMetadata.compute_num_computed_tokens@v0.30.0) + context_lens = common_attn_metadata.compute_num_computed_tokens().cpu() block_table = common_attn_metadata.block_table_tensor slot_mapping = common_attn_metadata.slot_mapping @@ -636,3 +653,31 @@ def wrapper(*args, **kwargs): return func(stack, *args, **kwargs) return wrapper + + +def framework_kv_block_size(backend_cls, *, cache_backend_cls=None, user_block_size: int | None = None) -> int: + """The KV page the FRAMEWORK would run this backend with — asked, not tabled. + + Mirrors serving (vllm 0.30.0): the cache manager block size is + ``get_preferred_block_size(CacheConfig.DEFAULT_BLOCK_SIZE)`` of the model's + FIRST non-SSM attention layer unless the user passed --block-size + (platforms/interface.py:588-640, _find_non_ssm_backend + + update_block_size_for_backend phase 1) — ``cache_backend_cls``, defaulting to + ``backend_cls`` when the measured layer is that first layer — and the kernel + page is ``select_common_block_size(cache_block_size, [backend])`` + (v1/worker/utils.py: the manager size when every backend accepts it, else the + largest explicit size that divides it). The former per-SM table (removed) was the owner's + 2026-09-30 interim; the B200 re-run on 2026-10-03 showed it is not + equivalent: with mla=32 the bf16 DSA cell selected the TRT-LLM FMHA sparse + cubin while serving ran FA-cute (page 64) — the page is a per-(backend, kv + dtype, head shape) decision, so it must come from the backend the case + actually selected. + """ + from vllm.config.cache import CacheConfig + from vllm.v1.worker.utils import select_common_block_size + + cache_backend = cache_backend_cls or backend_cls + cache_block = ( + user_block_size if user_block_size else cache_backend.get_preferred_block_size(CacheConfig.DEFAULT_BLOCK_SIZE) + ) + return int(select_common_block_size(int(cache_block), [backend_cls])) diff --git a/python/aisimulate/collector/wideep/trtllm/collect_moe_compute.py b/python/aisimulate/collector/wideep/trtllm/collect_moe_compute.py index f9d48fea64..673981452d 100644 --- a/python/aisimulate/collector/wideep/trtllm/collect_moe_compute.py +++ b/python/aisimulate/collector/wideep/trtllm/collect_moe_compute.py @@ -897,8 +897,8 @@ def run_moe_ep( # Re-verify on the next framework version bump. if moe_type == "fp8_block" and get_sm_version() != 90: raise ValueError( - f"WideEP fp8_block grouped GEMM is DeepGEMM-backed and SM90-only in " - f"trtllm 1.3.0rc20 (deep_gemm/compiler.cuh:330); got SM{get_sm_version()}" + f"FIXME(kernel-limit): WideEP fp8_block grouped GEMM is DeepGEMM-backed and SM90-only " + f"(deep_gemm/compiler.cuh:330 @1.3.0rc20, unchanged @1.3.0rc29); got SM{get_sm_version()}" ) _install_slot_lora_pad_patch() @@ -1210,29 +1210,27 @@ def run_forward_pass(): moe_ref.forward(hidden_states, actual_logits, do_finalize=not min_latency_mode) # ===================================================================== - # Benchmark with automatic power measurement and graph fallback - # (same pattern as collect_moe.py) + # Benchmark under CUDA-graph replay with automatic power measurement + # (same contract as collect_moe.py: a failed capture fails the case) # ===================================================================== num_warmups = 1 if distributed == "power_law" else 3 num_runs = 1 if distributed == "power_law" else 6 + # CUDA-graph timing is mandatory: a failed capture raises + # (benchmark_with_power default allow_graph_fail=False) and the case + # fails classified; the moe_ep row has no column that could record an + # eager-timed measurement. with benchmark_with_power( device=device, kernel_func=run_forward_pass, num_warmups=num_warmups, num_runs=num_runs, repeat_n=1, - allow_graph_fail=True, ) as results: # Calculate per-iteration latency (accounting for internal iterations) latency = results["latency_ms"] / num_iter power_stats = results["power_stats"] - # Always print CUDA graph status for debugging - print(f"[DEBUG] used_cuda_graph={results['used_cuda_graph']}, latency_ms={results['latency_ms']:.4f}") - if not results["used_cuda_graph"]: - print(f"CUDA graph capture failed for {num_tokens} tokens, used eager execution fallback") - # Determine source name if min_latency_mode: source = f"wideep_compute_{moe_kernel}_min_latency" diff --git a/python/aisimulate/src/aisimulate/generator/config/backend_config_mapping.yaml b/python/aisimulate/src/aisimulate/generator/config/backend_config_mapping.yaml index 49ccea4374..d287fe2726 100644 --- a/python/aisimulate/src/aisimulate/generator/config/backend_config_mapping.yaml +++ b/python/aisimulate/src/aisimulate/generator/config/backend_config_mapping.yaml @@ -169,16 +169,25 @@ parameters: sglang: max-running-requests trtllm: max_batch_size +# Token budgets render as multiples of 64 on EVERY backend (one rule, here, +# never per formula): engines size per-cache-group slot-mapping rows and +# scheduler chunks by these budgets, and a non-idiomatic width misaligns +# every row after the first (vLLM 0.29 DSV4: 6012 crashes the cutedsl +# compress kernel at first forward, 6016 passes). Rounding is always UP — +# a lower budget can violate scheduler minimums. Owner decision 2026-09-24: +# align all budgets uniformly instead of case by case. - param_key: enable_chunked_prefill vllm: enable-chunked-prefill - sglang: + sglang: key: chunked-prefill-size - value: "max_num_tokens if max_num_tokens else -1" + value: "((max_num_tokens + 63) // 64) * 64 if max_num_tokens else -1" trtllm: enable_chunked_prefill - param_key: max_prefill_tokens vllm: null - sglang: max-prefill-tokens + sglang: + key: max-prefill-tokens + value: "((max_prefill_tokens + 63) // 64) * 64 if max_prefill_tokens is defined and max_prefill_tokens is not none else none" trtllm: null - param_key: enable_mixed_chunk @@ -187,9 +196,16 @@ parameters: trtllm: null - param_key: max_num_tokens - vllm: max-num-batched-tokens + # same 64-alignment rule as above (the vLLM 0.29 DSV4 slot-mapping A/B is + # where it was found: 6012 crashes, 6016 passes); trtllm's block rounding + # in trtllm.rule (tokens_per_block = 32) is compatible with it. + vllm: + key: max-num-batched-tokens + value: "((max_num_tokens + 63) // 64) * 64 if max_num_tokens is defined and max_num_tokens is not none else none" sglang: null - trtllm: max_num_tokens + trtllm: + key: max_num_tokens + value: "((max_num_tokens + 63) // 64) * 64 if max_num_tokens is defined and max_num_tokens is not none else none" - param_key: max_seq_len vllm: max-model-len @@ -261,6 +277,22 @@ parameters: sglang: null trtllm: sparse_attention_config.implementation +# MiniMax-M3 sparse-attention algorithm key (trtllm >= 1.3.0rc29 requires it to build +# the M3 KV-cache manager; rendered by the rc29 engine template; set by utils.msa_sparse_algorithm) +- param_key: msa_sparse_algorithm + vllm: null + sglang: null + trtllm: sparse_attention_config.algorithm + +# sglang >= 0.5.20 prefill CUDA-graph backend ({full,breakable,tc_piecewise,disabled}). +# Rendered only when a deployment sets it: `breakable` is already the CUDA default +# (runner/prefill_cuda_graph_runner.py:17 @0.5.21), which mixed-chunk prefill requires +# (:291-296). CLI only on sglang. +- param_key: cuda_graph_backend_prefill + vllm: null + sglang: cuda-graph-backend-prefill + trtllm: null + - param_key: disable_overlap_scheduler vllm: null sglang: disable-overlap-schedule diff --git a/python/aisimulate/src/aisimulate/generator/config/backend_templates/sglang/cli_args.0.5.20.j2 b/python/aisimulate/src/aisimulate/generator/config/backend_templates/sglang/cli_args.0.5.20.j2 index aa897c9285..f58afcb5b3 100644 --- a/python/aisimulate/src/aisimulate/generator/config/backend_templates/sglang/cli_args.0.5.20.j2 +++ b/python/aisimulate/src/aisimulate/generator/config/backend_templates/sglang/cli_args.0.5.20.j2 @@ -1,6 +1,14 @@ {# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 +SGLang >= 0.5.20: ServerArgs became a msgspec Struct whose argparse flags are +derived from field names, and the old spellings `--cuda-graph-bs` / +`--cuda-graph-max-bs` (aliases of the *-decode fields in 0.5.16-0.5.19) are +gone: 0.5.21 rejects `--cuda-graph-bs` as ambiguous (matches +--cuda-graph-bs-decode / --cuda-graph-bs-prefill). The unified param keys are +unchanged; only the rendered spelling differs from cli_args.0.5.11.j2. +Verified against lmsysorg/sglang:v0.5.21 `launch_server --help` (2026-10-01); +boundary from v0.5.20 server_args.py (Struct form, no alias). -#} {%- set args = [] -%} {%- if sglang['tensor-parallel-size'] is defined and sglang['tensor-parallel-size'] is not none -%}{%- set _ = args.append('--tensor-parallel-size "' ~ sglang['tensor-parallel-size'] ~ '"') -%}{%- endif -%} @@ -48,4 +56,5 @@ SPDX-License-Identifier: Apache-2.0 {%- if sglang['speculative-num-draft-tokens'] is defined and sglang['speculative-num-draft-tokens'] is not none -%}{%- set _ = args.append('--speculative-num-draft-tokens "' ~ sglang['speculative-num-draft-tokens'] ~ '"') -%}{%- endif -%} {%- if sglang['disable-overlap-schedule'] -%}{%- set _ = args.append('--disable-overlap-schedule') -%}{%- endif -%} {%- if sglang['load-balance-method'] is defined and sglang['load-balance-method'] is not none -%}{%- set _ = args.append('--load-balance-method "' ~ sglang['load-balance-method'] ~ '"') -%}{%- endif -%} +{%- if sglang['cuda-graph-backend-prefill'] is defined and sglang['cuda-graph-backend-prefill'] is not none -%}{%- set _ = args.append('--cuda-graph-backend-prefill "' ~ sglang['cuda-graph-backend-prefill'] ~ '"') -%}{%- endif -%} {{ args | join(' ') }} diff --git a/python/aisimulate/src/aisimulate/generator/config/backend_templates/trtllm/extra_engine_args.1.3.0rc29.yaml.j2 b/python/aisimulate/src/aisimulate/generator/config/backend_templates/trtllm/extra_engine_args.1.3.0rc29.yaml.j2 new file mode 100644 index 0000000000..088fa110d3 --- /dev/null +++ b/python/aisimulate/src/aisimulate/generator/config/backend_templates/trtllm/extra_engine_args.1.3.0rc29.yaml.j2 @@ -0,0 +1,88 @@ +{# +SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +TRT-LLM >= 1.3.0rc29: copy of the rc23 template plus sparse_attention_config.algorithm +(MiniMax-M3 requires algorithm=minimax_m3 to build its KV-cache manager; op-probe harness +A/B 2026-10-01, facts/upgrade_rc29/ab_m3). Prior templates untouched. +-#} +backend: pytorch + +{% if is_moe is defined %} # is_moe from sdk +moe_expert_parallel_size: {{ moe_expert_parallel_size }} +moe_tensor_parallel_size: {{ moe_tensor_parallel_size }} + +moe_config: +{# moe_config.backend (facts / --generator-set params.agg.moe_config.backend) wins; the generic `moe_backend` param + (backend_config_mapping.yaml, the sglang --moe-runner-backend spelling) used to be silently dropped for trtllm — + it now feeds the same field, upper-cased to the TRT-LLM enum. Unset renders AUTO, rc29's own MoeConfig default + (llmapi/llm_args.py:1657-1659): an explicit CUTLASS fails DeepSeek FP8_BLOCK_SCALES on SM100 + (findings.yaml sm100_moe_backend_fact_rejected_2026_10_04), and rc29 no longer accepts WIDEEP. #} + backend: {{ moe_config.backend | default((moe_backend | upper) if moe_backend is defined and moe_backend else 'AUTO') }} # MOE backend to use (AUTO, CUTLASS, CUTEDSL, TRTLLM, DEEPGEMM, VANILLA, MARLIN, ...) + {% if moe_config.load_balancer is defined %} # Configuration for MoE load balancing + load_balancer: {{ moe_config.load_balancer }} # moe load balancer + {% endif %} +{% endif %} + +tensor_parallel_size: {{ tensor_parallel_size }} +pipeline_parallel_size: {{ pipeline_parallel_size }} +enable_attention_dp: {{ enable_attention_dp | default(false) }} +enable_chunked_prefill: {{ enable_chunked_prefill | default(false) }} + +{% set _max_seq_len = max_seq_len | default(none) %} + +max_batch_size: {{ max_batch_size }} +max_num_tokens: {{ max_num_tokens }} +{% if _max_seq_len is not none and _max_seq_len != '' %} +max_seq_len: {{ _max_seq_len }} +{% endif %} + +kv_cache_config: + {% if kv_cache_config.free_gpu_memory_fraction is defined %} + free_gpu_memory_fraction: {{ kv_cache_config.free_gpu_memory_fraction }} # Fraction of GPU memory usable for KV-cache + {% elif kv_cache_config.max_tokens is not defined %} + free_gpu_memory_fraction: 0.80 # Fraction of GPU memory usable for KV-cache + {% endif %} + {% if kv_cache_config.max_tokens is defined %} + max_tokens: {{ kv_cache_config.max_tokens }} + {% endif %} + dtype: {{ kv_cache_config.dtype | default('auto') }} + {% if kv_cache_config.tokens_per_block is defined %} + tokens_per_block: {{ kv_cache_config.tokens_per_block }} + {% endif %} + enable_block_reuse: {{ kv_cache_config.enable_block_reuse | default(false) }} + +{% if cache_transceiver_config.max_tokens_in_buffer is defined or cache_transceiver_config.backend is defined %} +cache_transceiver_config: + backend: {{ cache_transceiver_config.backend | default('DEFAULT') }} # The communication backend type to use for the cache transceiver ("DEFAULT","UCX","NIXL","MPI") +{% if cache_transceiver_config.max_tokens_in_buffer is defined %} + max_tokens_in_buffer: {{ cache_transceiver_config.max_tokens_in_buffer }} +{% endif %} +{% endif %} + +cuda_graph_config: + enable_padding: {{ cuda_graph_config.enable_padding | default(false) }} # Pad CUDA graph + batch_sizes: [{% for s in cuda_graph_config.batch_sizes + | default([1,2,4,8,16,32,64,128,256]) -%} + {{- s }}{{ ',' if not loop.last else '' }} {%- endfor %}] + +disable_overlap_scheduler: {{ disable_overlap_scheduler | default(true) }} +print_iter_log: {{ print_iter_log | default(false) }} + +{% if speculative_config.decoding_type is defined %} # speculative decoding +speculative_config: + decoding_type: {{ speculative_config.decoding_type }} + {% if speculative_config.decoding_type == 'MTP' and speculative_config.num_nextn_predict_layers is defined %} + num_nextn_predict_layers: {{ speculative_config.num_nextn_predict_layers }} + {% endif %} +{% endif %} + +{% if msa_sparse_algorithm is defined or msa_sparse_implementation is defined %} # MiniMax-M3 MSA +sparse_attention_config: +{% if msa_sparse_algorithm is defined %} + algorithm: {{ msa_sparse_algorithm }} # required by rc29+ (modeling_minimaxm3.py:337) on every SM +{% endif %} +{% if msa_sparse_implementation is defined %} + implementation: {{ msa_sparse_implementation }} # SM100-family: run the fmha_sm100 path the perf data represents +{% endif %} +{% endif %} diff --git a/python/aisimulate/src/aisimulate/generator/facts/apply.py b/python/aisimulate/src/aisimulate/generator/facts/apply.py index 8818a148a5..7a165255c4 100644 --- a/python/aisimulate/src/aisimulate/generator/facts/apply.py +++ b/python/aisimulate/src/aisimulate/generator/facts/apply.py @@ -96,7 +96,66 @@ def apply_model_default_args( tokens.append(_stringify(value)) -def apply_moe_backend(context: dict[str, Any], hardware: dict[str, Any] | None, *, backend: str) -> None: +def quant_family_of(quantization: dict[str, Any] | None) -> str | None: + """Collapse an artifact's quantization block (config.json ``quantization_config`` merged with the + modelopt ``hf_quant_config``, see utils._bundled_quantization) onto the collector's quant families: + nvfp4 / mxfp4 / fp8_block / fp8 / int4_wo. None when the artifact declares no quantization.""" + if not quantization: + return None + algo = str(quantization.get("quant_algo") or "").upper() + method = str(quantization.get("quant_method") or "").lower() + if algo == "MIXED_PRECISION": + # modelopt per-layer table: the MoE experts' algo decides the MoE kernel family (nvidia/Qwen3.8-2.4T-A95B-NVFP4: + # 92 `*.mlp.experts` layers NVFP4, 506 linear-attention / dense layers FP8); fall back to the most common algo + layers = quantization.get("quantized_layers") or {} + algos = [str((v or {}).get("quant_algo") or "").upper() for k, v in layers.items() if isinstance(v, dict)] + experts = [ + str((v or {}).get("quant_algo") or "").upper() + for k, v in layers.items() + if isinstance(v, dict) and "experts" in str(k) + ] + pick = experts or algos + algo = max(set(pick), key=pick.count) if pick else "" + quantization = {**quantization, "quant_algo": algo} + if algo == "NVFP4" or method in ("modelopt_fp4", "nvfp4"): + return "nvfp4" + if algo == "MXFP4" or method == "mxfp4": + return "mxfp4" + if algo.startswith("FP8") or method in ("fp8", "modelopt_fp8", "modelopt"): + return "fp8_block" if quantization.get("weight_block_size") else "fp8" + if ( + method in ("awq", "gptq", "compressed-tensors") + and int(quantization.get("bits") or quantization.get("weight_bits") or 0) == 4 + ): + return "int4_wo" + return method or algo.lower() or None + + +# TRT-LLM MoeConfig.backend literals a release no longer accepts, with the first release that rejects +# them: WIDEEP is in the literal through 1.3.0rc24 and gone from 1.3.0rc25 (llmapi/llm_args.py MoeConfig). +_TRTLLM_MOE_BACKEND_REJECTED_FROM = {"WIDEEP": "1.3.0rc25"} + + +def _trtllm_rejects_moe_backend(choice: str, backend_version: str | None) -> bool: + from packaging.version import InvalidVersion, Version + + rejected_from = _TRTLLM_MOE_BACKEND_REJECTED_FROM.get(str(choice).upper()) + if rejected_from is None or not backend_version: + return False + try: + return Version(str(backend_version)) >= Version(rejected_from) + except InvalidVersion: + return False + + +def apply_moe_backend( + context: dict[str, Any], + hardware: dict[str, Any] | None, + *, + backend: str, + quantization: dict[str, Any] | None = None, + backend_version: str | None = None, +) -> None: """Apply the hardware-derived ``moe_backend`` selection (fill-if-absent, MoE-only). ``hardware`` is the resolved hardware-profile dict (``ResolvedFacts.hardware``). @@ -112,21 +171,37 @@ def apply_moe_backend(context: dict[str, Any], hardware: dict[str, Any] | None, Per-backend seam: * ``trtllm``: sets ``context["moe_config"]["backend"]`` (consumed by the typed - engine-config builder). Only when unset — the builder defaults to CUTLASS. + engine-config builder). Only when unset — the template default applies otherwise + (AUTO from 1.3.0rc29, CUTLASS before). A literal the target ``backend_version`` + rejects (WIDEEP from 1.3.0rc25) is not filled. * ``sglang``: sets ``context["moe_backend"]`` (mapped to ``--moe-runner-backend`` via ``backend_config_mapping.yaml``). Only when unset. * ``vllm``: no-op — vLLM's MoE backend comes from MODEL ``defaults:`` (handled by :func:`apply_model_default_args`), not this hardware fact. """ choice = ((hardware or {}).get("moe_backend") or {}).get(backend) + if isinstance(choice, dict): + # keyed by the artifact's quantization family: on SM89 trtllm 1.3.0rc29 serves NVFP4 MoE only + # through MARLIN (weight-only), which in turn rejects bf16/fp8 MoE layers, so one hardware-wide + # value cannot be right (opharness findings sm89_trtllm_rc29_probe_2026_10_04). Unknown family + # -> the profile's "default" entry, else nothing (the builder/template default applies). + choice = choice.get(quant_family_of(quantization) or "", choice.get("default")) if not choice: return # GUARD: hardware moe_backend applies to MoE deployments only. if not context.get("is_moe"): return if backend == "trtllm": + if _trtllm_rejects_moe_backend(choice, backend_version): + return moe = context.setdefault("moe_config", {}) - if isinstance(moe, dict) and not moe.get("backend"): + # a user/recipe `moe_backend` param (the generic spelling the sglang branch below honours) counts as + # "already set" here too: until 2026-10-04 the hardware fill ignored it and the trtllm template never + # read it, so `--generator-set params.agg.moe_backend=...` was silently dropped on trtllm + user_set = context.get("moe_backend") or any( + context.get(f"{role}_moe_backend") for role in ("prefill", "decode", "agg", "encode") + ) + if isinstance(moe, dict) and not moe.get("backend") and not user_set: moe["backend"] = choice elif backend == "sglang": if not context.get("moe_backend"): diff --git a/python/aisimulate/src/aisimulate/generator/facts/hardware.yaml b/python/aisimulate/src/aisimulate/generator/facts/hardware.yaml index b193550d73..116282dcce 100644 --- a/python/aisimulate/src/aisimulate/generator/facts/hardware.yaml +++ b/python/aisimulate/src/aisimulate/generator/facts/hardware.yaml @@ -106,7 +106,14 @@ profiles: # SGLang: deepep_moe is the default wideep-compatible MoE runner on Blackwell. # Source: cli/main.py _sglang_moe_backend_override and # cli/exps/deepseek_disagg_sglang.yaml. - sglang: deepep_moe + # sglang: deepep_moe (not filled: no SGLang 0.5.x --moe-runner-backend choice) + # 2026-10-04 (B200 probe, PR #228 sm100 round): TRT-LLM 1.3.0rc29 rejects `WIDEEP` + # (pydantic literal_error; the literal is gone from 1.3.0rc25) and sglang 0.5.21 rejects + # `--moe-runner-backend deepep_moe` (argparse invalid choice; no 0.5.x lists it) — 156 + # trtllm + 150 sglang MoE renders failed to boot. facts/apply.py therefore drops WIDEEP + # for TRT-LLM >= 1.3.0rc25, where the rc29 template renders the framework's AUTO + # (A/B: explicit CUTLASS fails DeepSeek FP8_BLOCK_SCALES on SM100); older releases keep it. + # See benchmarks/evidence/collector/op-probe-2026-10/findings.yaml sm100_moe_backend_fact_rejected_2026_10_04. # --------------------------------------------------------------------------- # GB200 NVL72 — Blackwell sm100, arm64 Grace-Blackwell @@ -159,7 +166,14 @@ profiles: # SGLang: deepep_moe for the wideep DeepEP path on GB200 NVLink fabric. # Source: cli/main.py _sglang_moe_backend_override and # cli/exps/deepseek_disagg_sglang.yaml. - sglang: deepep_moe + # sglang: deepep_moe (not filled: no SGLang 0.5.x --moe-runner-backend choice) + # 2026-10-04 (B200 probe, PR #228 sm100 round): TRT-LLM 1.3.0rc29 rejects `WIDEEP` + # (pydantic literal_error; the literal is gone from 1.3.0rc25) and sglang 0.5.21 rejects + # `--moe-runner-backend deepep_moe` (argparse invalid choice; no 0.5.x lists it) — 156 + # trtllm + 150 sglang MoE renders failed to boot. facts/apply.py therefore drops WIDEEP + # for TRT-LLM >= 1.3.0rc25, where the rc29 template renders the framework's AUTO + # (A/B: explicit CUTLASS fails DeepSeek FP8_BLOCK_SCALES on SM100); older releases keep it. + # See benchmarks/evidence/collector/op-probe-2026-10/findings.yaml sm100_moe_backend_fact_rejected_2026_10_04. # --------------------------------------------------------------------------- # Intel Arc Pro B60 — XPU @@ -190,3 +204,76 @@ profiles: effect: NoSchedule moe_backend: trtllm: CUTLASS + + # --------------------------------------------------------------------------- + # RTX PRO Blackwell (PCIe, 8-GPU hosts) — sm120 + # Keyed by SM: the SDK carries several RTX PRO systems with the same hardware + # facts (rtx_pro_6000_server today; the RTX PRO 5000 campaign box will follow). + # Before this profile existed, resolve_facts() raised for every RTX render and + # the generator pipeline swallowed it, so NO model default (block-size, + # trust-remote-code, autotune flags) reached sm120 engine args + # (benchmarks/evidence/collector/op-probe-2026-10/findings.yaml: + # rtx_systems_have_no_generator_facts_hardware_profile, 2026-10-01). + # --------------------------------------------------------------------------- + sm120: + node_selector: + # GPU product label pattern NVIDIA-${name// /-} (same derivation as b200); + # device string observed in systems/data/rtx_pro_6000_server/*/*_perf.parquet + # `device`: "NVIDIA RTX PRO 6000 Blackwell Server Edition". + nvidia.com/gpu.product: NVIDIA-RTX-PRO-6000-Blackwell-Server-Edition + arch: x86_64 + # nccl_env: intentionally none. NCCL_CUMEM_ENABLE is unverified on sm120 and is + # not asserted. The sm120 campaign host routes cross-socket ring hops over + # NET/Socket by default (PCIe-only, two NUMA domains); the right interface + # setting is deployment-specific, not a hardware fact. + shared_memory: + # vLLM TP8 serving of GLM-5.2-NVFP4 on the 8x RTX PRO 5000 host ran with + # --shm-size 32g (2026-10-01, op-path vs forward-pass alignment run). + default: 32Gi + tolerations: + - key: nvidia.com/gpu + operator: Exists + effect: NoSchedule + # moe_backend: intentionally unset. trtllm WIDEEP needs the sm100 wide-EP + # fabric (PCIe sm120 has none; the typed builder defaults to CUTLASS when + # unset) and sglang deepep_moe is a wide-EP selection. vLLM picks its MoE + # backend from MODEL defaults (observed on sm120: FLASHINFER_CUTLASS NvFp4). + + # --------------------------------------------------------------------------- + # L40S — Ada sm89 (SDK system `l40s`) + # Product-keyed like h100 / h200 / h20 (the SM-keyed sm120 profile exists because several + # RTX PRO products share one set of facts; Ada does not — an L4 host would carry a + # different node_selector, so it gets its own profile when a campaign needs it). + # Before this profile existed resolve_facts() raised "Unknown hardware profile l40s" for + # every l40s render and the generator pipeline swallowed it ("Fact resolution failed; + # continuing without facts."), so no MODEL default (block-size, trust-remote-code, + # autotune flags) reached the engine args — the failure first found on sm120 + # (findings: rtx_systems_have_no_generator_facts_hardware_profile; sm89 campaign 2026-10-04). + # --------------------------------------------------------------------------- + l40s: + node_selector: + # GPU product label pattern NVIDIA-${name// /-} (same derivation as b200); the SDK's + # l40s perf data carries device "NVIDIA L40S". (The 2026-10-04 opharness campaign box + # was an 8x L40 — same AD102 die, lower TDP — which the system id does not cover.) + nvidia.com/gpu.product: NVIDIA-L40S + arch: x86_64 + # nccl_env: intentionally none (PCIe-only host, no NVLink; nothing verified to assert). + shared_memory: + # Not measured: every opharness probe on the L40 box (tp1 and tp4) ran with --shm-size 16g; + # 32Gi mirrors the other PCIe 8-GPU profile (sm120) as a conservative deployment default. + default: 32Gi + tolerations: + - key: nvidia.com/gpu + operator: Exists + effect: NoSchedule + # moe_backend: keyed by the artifact's quantization family (facts/apply.quant_family_of) because + # one hardware-wide value cannot be right on Ada: trtllm 1.3.0rc29 serves NVFP4 MoE only through + # MARLIN (weight-only; A/B on an L40: AUTO resolves to CUTLASS, which refuses NVFP4 below SM100, + # TRITON is SM90-only, DEEPGEMM needs SM100/103/107, VANILLA dies in CUDA-graph capture), while + # MARLIN rejects bf16 / fp8 MoE layers that CUTLASS serves. FP8-block and MXFP4 MoE have no + # implementation on SM89 at all, so no key is given for them (the template default CUTLASS + # fails exactly as every alternative does). No wide-EP fabric on PCIe Ada: no WIDEEP / + # deepep_moe. Evidence: opharness findings sm89_trtllm_rc29_probe_2026_10_04. + moe_backend: + trtllm: + nvfp4: MARLIN diff --git a/python/aisimulate/src/aisimulate/generator/facts/models.yaml b/python/aisimulate/src/aisimulate/generator/facts/models.yaml index c2a622724c..810f6e20fd 100644 --- a/python/aisimulate/src/aisimulate/generator/facts/models.yaml +++ b/python/aisimulate/src/aisimulate/generator/facts/models.yaml @@ -11,9 +11,11 @@ models: default_variant: pro default_precision: fp8 match: + # `DeepSeek-V4-*` (not `V4*`): DeepSeek-V4.1 is a different model with its + # own profile below; the broader glob used to hand it V4's defaults. model_paths: - - deepseek-ai/DeepSeek-V4* - - sgl-project/DeepSeek-V4* + - deepseek-ai/DeepSeek-V4-* + - sgl-project/DeepSeek-V4-* architectures: - DeepseekV4ForCausalLM model_types: @@ -113,3 +115,105 @@ models: job_name: model-download pvc_name: model-cache mount_path: /model-store + + minimax-m3: + id: minimax-m3 + match: + model_paths: + - MiniMaxAI/MiniMax-M3* + - nvidia/MiniMax-M3* + architectures: + - MiniMaxM3SparseForConditionalGeneration + - MiniMaxM3VLForConditionalGeneration + model_types: + - minimax_m3 + traits: + - moe + - sparse_attention + defaults: + # Provenance: op-probe harness on L40, vllm 0.30.0, 2026-10-05 + # (benchmarks/evidence/collector/op-probe-2026-10/findings.yaml + # minimax_m3_dummy_dense_head_2026_10_05). M3 is three dense-attention layers followed by sparse layers; vllm + # sizes the KV manager block to 16 by default, the dense layers' FlashAttention backend accepts MultipleOf(16) + # but MiniMaxM3SparseBackend / MiniMaxM3IndexerBackend accept exactly [128] (one sparse block per KV page, + # models/minimax_m3/common/sparse_attention.py), so select_common_block_size has no candidate and the engine + # dies at "No common block size for 16" (v1/worker/utils.py:386). 128 satisfies every layer's backend. An + # all-sparse dummy cut never showed this (its only backend forces 128), which is how the sm90 matrix passed. + - match: {backend: vllm} + roles: ["*"] + backend_args: + block-size: 128 + deepseek-v4.1: + id: deepseek-v4.1 + default_variant: flash + default_precision: fp8 + match: + model_paths: + - deepseek-ai/DeepSeek-V4.1* + - nvidia/DeepSeek-V4.1* + architectures: + - DeepseekV41ForCausalLM + model_types: + - deepseek_v41 + traits: + - moe + - mla + - mtp + - compressed_attention + - engram + requires: + capabilities: + - deepseek_v4_openai_parsing + - deepseek_v4_tokenization + defaults: + # Provenance: op-probe harness A/B on H20, vllm 0.30.0, 2026-10-01 + # (benchmarks/evidence/collector/op-probe-2026-10/findings.yaml + # dsv41_dummy_cut_and_vllm_block_size_2026_10_01). V4.1 inherits V4's + # trust-remote-code / autotune defaults but NOT block-size 256: its + # compressed layers (compress_ratio 1/2) run 64-token kernel blocks, and + # vllm 0.30.0 refuses to split a 256-token manager block 4-way under the + # BLHNC layout once a KV tensor is shared by several layers + # ("block stride != page", kv_cache_interface.py:373). 64 loads; it is the + # value vllm's own error message prescribes. + - match: {backend: vllm} + roles: ["*"] + backend_args: + trust-remote-code: true + block-size: 64 + no-enable-flashinfer-autotune: true + # Same sglang defaults as deepseek-v4 (V4.1 code lives in the same + # deepseek_v4 model file on sglang main; no release supports it yet). + - match: {backend: sglang} + roles: ["*"] + backend_args: + trust-remote-code: true + disable-flashinfer-autotune: true + precision_variants: + fp8: + weight_dtype: fp8 + activation_dtype: fp8 + expert_dtype: fp4 + kv_cache_dtype_default: fp8 + requires: + capabilities: + - fp8_kv_cache + - deepseek_v4_attention_indexer + - fp4_moe_runner + variants: + flash: + id: flash + match: + model_paths: + - deepseek-ai/DeepSeek-V4.1-Flash* + - nvidia/DeepSeek-V4.1-Flash* + default_precision: fp8 + artifacts: + include: + model_cache_pvc: + name: model-cache + storage: 400Gi + storageClassName: your-storage-class-name + model_download_job: + job_name: model-download + pvc_name: model-cache + mount_path: /model-store diff --git a/python/aisimulate/src/aisimulate/generator/facts/request_resolution.py b/python/aisimulate/src/aisimulate/generator/facts/request_resolution.py index dc7a023ec9..e568123478 100644 --- a/python/aisimulate/src/aisimulate/generator/facts/request_resolution.py +++ b/python/aisimulate/src/aisimulate/generator/facts/request_resolution.py @@ -42,6 +42,13 @@ "gb200": "gb200", "gb200_sxm": "gb200", "gb300": "gb200", + # RTX PRO Blackwell PCIe hosts share one SM-keyed profile; the RTX PRO 5000 + # system (SDK entry pending) maps here too once it exists. + "rtx_pro_6000_server": "sm120", + # Ada: product-keyed like the Hopper entries (the SDK has an l40s system and nothing else on + # sm89 yet; an l4 system would get its own profile). Explicit so the table stays the one place + # that says which systems have facts. + "l40s": "l40s", } # Sections of the params dict that may carry the SDK system id. diff --git a/python/aisimulate/src/aisimulate/generator/module_bridge.py b/python/aisimulate/src/aisimulate/generator/module_bridge.py index b68b494d5f..2a57f13112 100644 --- a/python/aisimulate/src/aisimulate/generator/module_bridge.py +++ b/python/aisimulate/src/aisimulate/generator/module_bridge.py @@ -28,6 +28,34 @@ def _msa_sparse_implementation(task_config) -> str | None: ) +def _msa_sparse_algorithm(task_config) -> str | None: + """Optimized-path wrapper over utils.msa_sparse_algorithm (shared with naive).""" + from .utils import msa_sparse_algorithm + + return msa_sparse_algorithm( + getattr(task_config, "primary_backend_name", None) or "", + task_config.primary_model_path, + ) + + +def _model_has_kda(model_path: str) -> bool: + from .utils import model_has_kda + + return model_has_kda(model_path) + + +def _vllm_dsa_kv_cache_dtype(task_config, gemm_quant_mode) -> str | None: + """Optimized-path wrapper over utils.vllm_dsa_kv_cache_dtype: NVFP4 DSA + checkpoints on vLLM must render an explicit fp8 kv-cache dtype.""" + from .utils import vllm_dsa_kv_cache_dtype + + return vllm_dsa_kv_cache_dtype( + getattr(task_config, "primary_backend_name", None) or "", + task_config.primary_model_path, + gemm_quant_mode, + ) + + def _deep_merge(target: dict, extra: dict | None) -> dict: """ Recursively merge the contents of the 'extra' dictionary into 'target', @@ -185,6 +213,11 @@ def _build_worker_params(prefix: str, extra_overrides: dict | None) -> tuple[dic worker_payload["memory"] = memory if quant.get("kvcache_quant_mode"): worker_payload["kv_cache_dtype"] = quant["kvcache_quant_mode"] + # NVFP4 DSA checkpoints pin the KV cache to FP8; vLLM's `auto` + # spelling of it is rejected by the sparse-MLA selector (utils docstring) + _dsa_kv = _vllm_dsa_kv_cache_dtype(task_config, quant.get("gemm_quant_mode")) + if _dsa_kv is not None: + worker_payload["kv_cache_dtype"] = _dsa_kv if encoder_dp is not None: worker_payload["enable_encoder_dp"] = encoder_dp if attention_backend is not None: @@ -270,6 +303,9 @@ def _build_worker_params(prefix: str, extra_overrides: dict | None) -> tuple[dic "nextn": task_config.nextn, "nextn_accepted": task_config.nextn_accepted if task_config.nextn else None, "msa_sparse_implementation": _msa_sparse_implementation(task_config), + "msa_sparse_algorithm": _msa_sparse_algorithm(task_config), + # KDA linear-attention layers: vllm.rule caps the decode batch (see there) + "has_kda": _model_has_kda(task_config.primary_model_path), } model_cfg = {k: v for k, v in model_cfg.items() if v is not None} model_cfg = _deep_merge(model_cfg, overrides.get("ModelConfig")) diff --git a/python/aisimulate/src/aisimulate/generator/naive.py b/python/aisimulate/src/aisimulate/generator/naive.py index 1ae8beb46c..edd2b8d5da 100644 --- a/python/aisimulate/src/aisimulate/generator/naive.py +++ b/python/aisimulate/src/aisimulate/generator/naive.py @@ -30,7 +30,7 @@ get_model_config_from_model_path, ) -from .utils import msa_sparse_implementation +from .utils import model_has_kda, msa_sparse_algorithm, msa_sparse_implementation, vllm_dsa_kv_cache_dtype logger = logging.getLogger(__name__) @@ -517,6 +517,12 @@ def build_naive_generator_params( "max_batch_size": max_batch_size, "gpus_per_worker": min_gpus, } + # NVFP4 DSA checkpoints on vLLM need an explicit fp8 kv-cache dtype (see + # utils.vllm_dsa_kv_cache_dtype); the naive path has no perf task, so the + # artifact fact comes from the bundled config + _kv_dtype = vllm_dsa_kv_cache_dtype(backend_name, model_name, model_config=model_config) + if _kv_dtype is not None: + worker_params["kv_cache_dtype"] = _kv_dtype name_prefix = _sanitize_rfc1123(model_name) @@ -548,6 +554,8 @@ def build_naive_generator_params( "is_moe": is_moe, "fits_in_memory": fits, "required_tp": required_tp, + # KDA linear-attention layers: vllm.rule caps the decode batch (see there) + "has_kda": model_has_kda(model_name, model_config), } # Shared MSA prescription (see utils.msa_sparse_implementation): the # naive entry point must emit the same MiniMax-M3/SM100-family @@ -557,6 +565,10 @@ def build_naive_generator_params( _msa_impl = msa_sparse_implementation(backend_name, model_name, system_name) if _msa_impl is not None: model_config["msa_sparse_implementation"] = _msa_impl + # TRT-LLM >= rc29 requires sparse_attention_config.algorithm for M3 (every SM) + _msa_algo = msa_sparse_algorithm(backend_name, model_name) + if _msa_algo is not None: + model_config["msa_sparse_algorithm"] = _msa_algo service = _deep_merge_dicts(service, _section_override(overrides, "ServiceConfig")) k8s = _deep_merge_dicts(k8s, _section_override(overrides, "K8sConfig")) diff --git a/python/aisimulate/src/aisimulate/generator/rendering/engine.py b/python/aisimulate/src/aisimulate/generator/rendering/engine.py index 3798b022aa..6c295110be 100644 --- a/python/aisimulate/src/aisimulate/generator/rendering/engine.py +++ b/python/aisimulate/src/aisimulate/generator/rendering/engine.py @@ -452,8 +452,17 @@ def render_backend_templates( # a model profile). No-op without resolved facts. if resolved_facts is not None and backend == "trtllm": from aisimulate.generator.facts.apply import apply_moe_backend - - apply_moe_backend(context, getattr(resolved_facts, "hardware", None), backend=backend) + from aisimulate.generator.utils import _bundled_quantization + + hardware = getattr(resolved_facts, "hardware", None) + # the artifact's quantization decides the MoE kernel family on SMs where one value cannot + # serve every quant (hardware.yaml l40s); None for unquantized / unresolvable checkpoints. + # Loaded only for that case: it may fetch the config from the Hub. + keyed = isinstance(((hardware or {}).get("moe_backend") or {}).get(backend), dict) + quantization = ( + _bundled_quantization(str(context.get("model_path") or "")) if keyed and context.get("is_moe") else None + ) + apply_moe_backend(context, hardware, backend=backend, quantization=quantization, backend_version=version) # Find template files template_path = Path(templates_dir) diff --git a/python/aisimulate/src/aisimulate/generator/rendering/rule_engine.py b/python/aisimulate/src/aisimulate/generator/rendering/rule_engine.py index f7fc5e06e3..c1faff255e 100644 --- a/python/aisimulate/src/aisimulate/generator/rendering/rule_engine.py +++ b/python/aisimulate/src/aisimulate/generator/rendering/rule_engine.py @@ -173,6 +173,48 @@ def _load_rule_path(base_dir: str, backend: str) -> Optional[str]: return p if os.path.exists(p) else None +_TOKEN_BUDGET_KEYS = ("max_num_tokens", "max_prefill_tokens") +_TOKEN_BUDGET_ALIGN = 64 + + +def _align_token_budgets(param_values: dict[str, Any]) -> None: + """Round every token budget UP to a multiple of 64, on every backend, in + every role — the ONE alignment rule (owner decision 2026-09-24). Engines + size per-cache-group slot-mapping rows and scheduler chunks by these + budgets and a non-idiomatic width misaligns every row after the first + (vLLM 0.29 DeepSeek-V4: 6012 crashes the cutedsl compress kernel, 6016 + passes). Applied here, after the rule formulas and before any rendering, + so the trtllm engine-yaml template — which the mapping transforms cannot + reach — gets it too; the mapping keeps the same rounding as a guard for + values injected past the rule engine. Never rounds down. Values a caller + already evaluated (``preserve_engine_limits``, the Sweeper) are rounded too, + with a warning, so the deployed limit is visible next to the simulated one.""" + params = param_values.get("params") + if not isinstance(params, dict): + return + preserved = bool(param_values.get("preserve_engine_limits") or params.get("preserve_engine_limits")) + for role_name, role in params.items(): + if not isinstance(role, dict): + continue + for key in _TOKEN_BUDGET_KEYS: + v = role.get(key) + if isinstance(v, bool) or not isinstance(v, int) or v <= 0: + continue + aligned = ((v + _TOKEN_BUDGET_ALIGN - 1) // _TOKEN_BUDGET_ALIGN) * _TOKEN_BUDGET_ALIGN + if aligned != v: + if preserved: + logger.warning( + "%s %s=%d was evaluated by the caller; the deployment uses %d (token budgets are " + "rounded up to a multiple of %d on every backend)", + role_name, + key, + v, + aligned, + _TOKEN_BUDGET_ALIGN, + ) + role[key] = aligned + + def apply_rule_plugins( param_values: dict[str, Any], backend: str, @@ -182,6 +224,7 @@ def apply_rule_plugins( base = _resolve_rule_plugin_dir(rule_name, base_dir=dsl_dir) rule_path = _load_rule_path(base, backend) if not rule_path: + _align_token_budgets(param_values) return param_values with open(rule_path, encoding="utf-8") as f: content = f.read().splitlines() @@ -227,4 +270,5 @@ def apply_rule_plugins( _apply_line(assign, backend, param_values, default_scope) except Exception: logger.exception("rule apply failed at line %d", idx) + _align_token_budgets(param_values) return param_values diff --git a/python/aisimulate/src/aisimulate/generator/request/sweeper.py b/python/aisimulate/src/aisimulate/generator/request/sweeper.py index 9cb2b4458e..6042a0482d 100644 --- a/python/aisimulate/src/aisimulate/generator/request/sweeper.py +++ b/python/aisimulate/src/aisimulate/generator/request/sweeper.py @@ -455,6 +455,18 @@ def from_sweeper_candidate( value = getattr(resolved_model_facts, key) if value is not None: raw[f"ModelConfig.{key}"] = value + if backend == "trtllm": + # The benchmark rule renders the MiniMax-M3 sparse-attention keys from ModelConfig as the + # default path does (module_bridge): TRT-LLM >= 1.3.0rc29 needs `algorithm` to build the M3 + # KV-cache manager on every SM; `implementation` is the SM100-family prescription. + from aisimulate.generator.utils import msa_sparse_algorithm, msa_sparse_implementation + + for key, value in ( + ("msa_sparse_algorithm", msa_sparse_algorithm(backend, model_path)), + ("msa_sparse_implementation", msa_sparse_implementation(backend, model_path, hardware_sku)), + ): + if value and f"ModelConfig.{key}" not in raw: + raw[f"ModelConfig.{key}"] = value request = GeneratorRequest( model=ModelSpec(model_path=model_path), diff --git a/python/aisimulate/src/aisimulate/generator/rule_plugin/benchmark/trtllm.rule b/python/aisimulate/src/aisimulate/generator/rule_plugin/benchmark/trtllm.rule index 55850ed7b7..11a5fa412a 100644 --- a/python/aisimulate/src/aisimulate/generator/rule_plugin/benchmark/trtllm.rule +++ b/python/aisimulate/src/aisimulate/generator/rule_plugin/benchmark/trtllm.rule @@ -53,3 +53,10 @@ when (ModelConfig.nextn or 0) > 0: when kv_cache_dtype in ("float16", "bfloat16"): kv_cache_dtype = "auto" + +# MiniMax-M3 sparse attention, as in the default rule: rc29 requires algorithm on every SM, +# implementation is the SM100-family prescription (utils.msa_sparse_algorithm / _implementation). +when (ModelConfig.msa_sparse_implementation or '') != '': + agg_prefill_decode msa_sparse_implementation = ModelConfig.msa_sparse_implementation +when (ModelConfig.msa_sparse_algorithm or '') != '': + agg_prefill_decode msa_sparse_algorithm = ModelConfig.msa_sparse_algorithm diff --git a/python/aisimulate/src/aisimulate/generator/rule_plugin/trtllm.rule b/python/aisimulate/src/aisimulate/generator/rule_plugin/trtllm.rule index 50626d4df6..904d4fd8ab 100644 --- a/python/aisimulate/src/aisimulate/generator/rule_plugin/trtllm.rule +++ b/python/aisimulate/src/aisimulate/generator/rule_plugin/trtllm.rule @@ -73,3 +73,9 @@ when kv_cache_dtype in ("float16", "bfloat16"): # sm_version 100/103; byte-safe everywhere else. when (ModelConfig.msa_sparse_implementation or '') != '': agg_prefill_decode msa_sparse_implementation = ModelConfig.msa_sparse_implementation + +# TRT-LLM >= 1.3.0rc29 requires sparse_attention_config.algorithm='minimax_m3' for MiniMax-M3 on +# every SM (modeling_minimaxm3.py:337); utils.msa_sparse_algorithm sets it architecture-keyed, +# the rc29 engine template renders it (older templates ignore the variable). +when (ModelConfig.msa_sparse_algorithm or '') != '': + agg_prefill_decode msa_sparse_algorithm = ModelConfig.msa_sparse_algorithm diff --git a/python/aisimulate/src/aisimulate/generator/rule_plugin/vllm.rule b/python/aisimulate/src/aisimulate/generator/rule_plugin/vllm.rule index 81e690b19f..c9f024f545 100644 --- a/python/aisimulate/src/aisimulate/generator/rule_plugin/vllm.rule +++ b/python/aisimulate/src/aisimulate/generator/rule_plugin/vllm.rule @@ -7,6 +7,15 @@ when not (preserve_engine_limits or false): agg_decode max_batch_size = (512 if (max_batch_size or 0) < 512 else (max_batch_size * 2)) prefill max_batch_size = (max_batch_size if max_batch_size else 1) +# KDA (Kimi Delta Attention) models on vLLM: the 512-batch decode CUDA graph +# faults at replay right after capture (vLLM 0.29.0 / SM90, Kimi-K3: 64, 256, +# 384, 448, 480, 496 pass; 512 fails; graphs off at 512 passes — opharness +# findings followups_after_full_reprobe_2026_09_24 + k3 bisection). Owner +# decision 2026-09-25: cap the decode batch at 256 for KDA models. Applied +# BEFORE the token budgets below so max_num_tokens follows the capped batch. +when (ModelConfig.has_kda or false) and not (preserve_engine_limits or false): + agg_decode max_batch_size = (256 if (max_batch_size or 0) > 256 else max_batch_size) + # Enforce vLLM MoE parallelism: vLLM uses unified TP for both attention and MoE, # so tensor_parallel_size must account for moe_tp × moe_ep. # Without this, MoE models get tp=1 and the full model is replicated per GPU, causing OOM. @@ -23,6 +32,9 @@ agg_prefill_decode gpus_per_worker = (tensor_parallel_size or 1) * (pipeline_par agg_prefill_decode enable_expert_parallel = ((moe_expert_parallel_size or 1) > 1) when not (preserve_engine_limits or false): + # plain budgets; the 64-alignment every backend needs is applied ONCE in + # backend_config_mapping.yaml (max_num_tokens / max_prefill_tokens / + # chunked-prefill-size), not per formula prefill max_num_tokens = (SlaConfig.isl or 0) + 1500 decode max_num_tokens = max_batch_size agg max_num_tokens = (max_batch_size or 0) + (SlaConfig.isl or 0) + 1500 diff --git a/python/aisimulate/src/aisimulate/generator/utils.py b/python/aisimulate/src/aisimulate/generator/utils.py index 8c0d7f90e8..acbf86f4e5 100644 --- a/python/aisimulate/src/aisimulate/generator/utils.py +++ b/python/aisimulate/src/aisimulate/generator/utils.py @@ -173,3 +173,146 @@ def msa_sparse_implementation(backend_name: str, model_path: str, system_name: s if int(spec.get("gpu", {}).get("sm_version", -1)) in (100, 103): return "msa" return None + + +_MSA_ARCHITECTURES = ("MiniMaxM3ForCausalLM", "MiniMaxM3SparseForConditionalGeneration") + + +def msa_sparse_algorithm(backend_name: str, model_path: str) -> str | None: + """MiniMax-M3 x TRT-LLM: the ``sparse_attention_config.algorithm`` key. + + TRT-LLM >= 1.3.0rc29 refuses to build the M3 KV-cache manager without + ``sparse_attention_config.algorithm='minimax_m3'`` (modeling_minimaxm3.py:337 + ``_validate_sparse_attention_runtime_config`` @1.3.0rc29) on every SM — unlike + the ``implementation`` knob above, which is an SM100-family performance + prescription. Keyed on the checkpoint ARCHITECTURE; None for every other + model/backend so the field is dropped. Rendered by the >= rc29 engine + template only (older templates never emitted the key and rc23 did not read it). + """ + if backend_name != "trtllm": + return None + from aisimulate.sdk.utils import get_model_config_from_model_path + + try: + architecture = get_model_config_from_model_path(model_path).get("architecture") + except Exception: # FileNotFound / KeyError / ValueError / HuggingFaceDownloadError + return None + return "minimax_m3" if architecture in _MSA_ARCHITECTURES else None + + +def _model_architecture(model_path: str, model_config: dict | None = None) -> str | None: + """Architecture of a checkpoint config, None when unresolvable (a + user-local or unreachable checkpoint the SDK cannot load — the render + then carries no prescription, it never fails on this fact). A frozen + ``model_config`` (naive path) is used as-is: no resolution runs. + Indirection so tests can stub the SDK lookup.""" + if isinstance(model_config, dict): + return model_config.get("architecture") + from aisimulate.sdk.utils import get_model_config_from_model_path + + try: + return get_model_config_from_model_path(model_path).get("architecture") + except Exception: # FileNotFound / KeyError / ValueError / HuggingFaceDownloadError + return None + + +# DeepSeek sparse attention (DSA) architectures: the sparse-MLA selector is +# the one that rejects vLLM's auto-resolved kv dtype spelling (see below). +_DSA_ARCHITECTURES = ("GlmMoeDsaForCausalLM", "DeepseekV32ForCausalLM") + + +def _bundled_quantization(model_path: str, raw: dict | None = None) -> dict | None: + """The artifact's quantization facts as the SDK loader exposes them: + config.json ``quantization_config`` merged with the ``hf_quant_config`` + the loader attaches from the bundled ``_hf_quant_config.json`` (the + modelopt artifacts keep ``kv_cache_quant_algo`` ONLY there — their + config.json has no quantization block at all). None when the config is + unresolvable or carries neither. Indirection so tests can stub it.""" + if raw is None: + from aisimulate_core.sdk.utils import _load_model_config_from_model_path + + try: + raw = _load_model_config_from_model_path(model_path) + except Exception: # unresolvable checkpoint: no quantization fact + return None + merged: dict = {} + hfq = raw.get("hf_quant_config") + if isinstance(hfq, dict): + merged.update(hfq.get("quantization") if isinstance(hfq.get("quantization"), dict) else hfq) + if isinstance(raw.get("quantization_config"), dict): + merged.update(raw["quantization_config"]) + return merged or None + + +def _artifact_pins_fp8_kv(quantization: dict | None) -> bool: + """modelopt artifacts pin the KV cache either as ``kv_cache_quant_algo: + FP8`` (hf_quant_config.json) or as ``kv_cache_scheme: {num_bits: 8, + type: float}`` (config.json quantization block); either spelling counts.""" + if not quantization: + return False + if str(quantization.get("kv_cache_quant_algo") or "").upper() == "FP8": + return True + scheme = quantization.get("kv_cache_scheme") + if isinstance(scheme, str): + return scheme.upper() == "FP8" + scheme = scheme or {} + return scheme.get("num_bits") == 8 and str(scheme.get("type", "")).lower() == "float" + + +def vllm_dsa_kv_cache_dtype( + backend_name: str, model_path: str, gemm_quant_mode: Any = None, model_config: dict | None = None +) -> str | None: + """NVFP4 DSA checkpoints x vLLM: prescribe ``--kv-cache-dtype fp8``. + + The modelopt NVFP4 artifacts of the DSA models (GLM-5 / 5.1 / 5.2 / 5.3 + -NVFP4, DeepSeek-V3.2-NVFP4) pin the KV cache to FP8. vLLM 0.29.0 + resolves ``--kv-cache-dtype auto`` for them to the literal ``fp8_e4m3`` + and its sparse-MLA backend selector rejects that spelling + (FLASHMLA_SPARSE: "kv_cache_dtype not supported"); the same engine with + an explicit ``fp8`` loads FlashMLASparseImpl and serves. Serving-side + the KV IS fp8 either way (checkpoint scheme), so the explicit value only + says what ``auto`` means for these artifacts. + + Keyed on the checkpoint ARCHITECTURE (DSA) plus an artifact fact — never + a model-name pattern: the task's ``nvfp4`` GEMM quant mode on the + optimized path, or the bundled config's fp8 KV scheme on the naive + ``cli generate`` path (which has no perf task). Bundled configs without a + ``quantization_config`` (the Hub config of DeepSeek-V3.2-NVFP4 keeps its + quantization only in hf_quant_config.json) get no prescription until the + SDK bundles that fact — a known gap, recorded in the opharness findings. + Evidence: findings ``vllm_fp8kv`` (0.29 addendum), + ``glm53_onboarding_2026_09_24``; owner decision 2026-09-24. Returns None + everywhere else so the task's own kv mode stands. + """ + if backend_name != "vllm": + return None + if _model_architecture(model_path, model_config) not in _DSA_ARCHITECTURES: + return None + if str(gemm_quant_mode or "").lower() == "nvfp4": + return "fp8" + raw = model_config.get("raw_config") if isinstance(model_config, dict) else None + if _artifact_pins_fp8_kv(_bundled_quantization(model_path, raw if isinstance(raw, dict) else None)): + return "fp8" + return None + + +def model_has_kda(model_path: str, model_config: dict | None = None) -> bool: + """True when the checkpoint declares KDA (Kimi Delta Attention) linear-attention + layers: config.json (or text_config) linear_attn_config.kda_layers non-empty + (Kimi-K3, Kimi-Linear, GLM-5.3-Flash). Read from the SDK-loaded config so + bundled, local and Hub checkpoints resolve identically; any failure to load + is False (the caller renders the default, nothing is invented).""" + if isinstance(model_config, dict): # frozen config (naive path): no resolution + raw = model_config.get("raw_config") or {} + else: + try: + from aisimulate.sdk.utils import get_model_config_from_model_path + + raw = get_model_config_from_model_path(model_path).get("raw_config") or {} + except Exception: # unresolvable checkpoint: no fact, no override + return False + for cfg in (raw.get("text_config") or {}, raw): + la = cfg.get("linear_attn_config") if isinstance(cfg, dict) else None + if isinstance(la, dict) and la.get("kda_layers"): + return True + return False diff --git a/python/aisimulate/src/aisimulate_core/sdk/common.py b/python/aisimulate/src/aisimulate_core/sdk/common.py index ff135bf849..0ddaaf7bc5 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/common.py +++ b/python/aisimulate/src/aisimulate_core/sdk/common.py @@ -1408,14 +1408,13 @@ class PerfDataFilename(Enum): # DeepSeek-V4 sparse-op family — all share one column schema, served by # the engine table view (``_dsv4_sparse_kernel_data.*`` / # ``_dsv4_csa_topk_calib_data``): - # csa_attn / hca_attn / paged_mqa_logits : FMLA & indexer kernel latency, - # keyed ``num_heads -> tp -> past_kv -> isl -> bs`` (kernel-level Δ data, - # queried by ``_lookup_sparse_kernel``). + # paged_mqa_logits : indexer kernel latency, keyed + # ``num_heads -> tp -> past_kv -> isl -> bs`` (kernel-level Δ data, + # queried by ``_lookup_sparse_kernel``). The csa_attn / hca_attn FMLA + # sidecars were removed 2026-10-15 (retired placeholders, no query path). # csa_topk_calib : two rows/shape (score_mode=flat|top_last); the topK # DELTA (flat-top_last) correction applied to CSA module latency. dsv4_paged_mqa_logits_module = "dsv4_paged_mqa_logits_module_perf.parquet" - dsv4_hca_attn_module = "dsv4_hca_attn_module_perf.parquet" - dsv4_csa_attn_module = "dsv4_csa_attn_module_perf.parquet" dsv4_csa_topk_calib = "dsv4_csa_topk_calib_perf.parquet" dsv4_megamoe_module = "dsv4_megamoe_module_perf.parquet" # Whole-model forward-pass data (forward_model="fpm"). Paired with a diff --git a/python/aisimulate/src/aisimulate_core/sdk/operations/dsv4.py b/python/aisimulate/src/aisimulate_core/sdk/operations/dsv4.py index c0c0dbe77c..6daae8b15e 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/operations/dsv4.py +++ b/python/aisimulate/src/aisimulate_core/sdk/operations/dsv4.py @@ -177,8 +177,10 @@ def _load_sparse(sub_key, filename_enum): # paged_mqa_logits is the only sparse-kernel sidecar with a live # consumer (the CP chunk walk). The hca_attn/csa_attn sidecars - # were design placeholders with no query path and were retired - # (csa_attn never even shipped data). + # were design placeholders with no query path; they were retired + # here first and removed end to end (collector ops, PerfFile / + # PerfDataFilename, shipped tables) on 2026-10-15 — csa_attn never + # even shipped data. sparse_loaded = { "paged_mqa_logits": _load_sparse("paged_mqa_logits", PerfDataFilename.dsv4_paged_mqa_logits_module), } diff --git a/python/aisimulate/src/aisimulate_core/systems/data/b200_sxm/sparse_attention/vllm/0.24.0/collection_meta.yaml b/python/aisimulate/src/aisimulate_core/systems/data/b200_sxm/sparse_attention/vllm/0.24.0/collection_meta.yaml index 43c67faf60..45041ba8f6 100644 --- a/python/aisimulate/src/aisimulate_core/systems/data/b200_sxm/sparse_attention/vllm/0.24.0/collection_meta.yaml +++ b/python/aisimulate/src/aisimulate_core/systems/data/b200_sxm/sparse_attention/vllm/0.24.0/collection_meta.yaml @@ -27,8 +27,6 @@ tables: status: complete dsv4_csa_generation_module_perf: status: complete - dsv4_hca_attn_module_perf: - status: complete dsv4_hca_context_module_perf: status: complete dsv4_hca_generation_module_perf: diff --git a/python/aisimulate/src/aisimulate_core/systems/data/b200_sxm/sparse_attention/vllm/0.24.0/dsv4_hca_attn_module_perf.parquet b/python/aisimulate/src/aisimulate_core/systems/data/b200_sxm/sparse_attention/vllm/0.24.0/dsv4_hca_attn_module_perf.parquet deleted file mode 100644 index f3333ffe05..0000000000 Binary files a/python/aisimulate/src/aisimulate_core/systems/data/b200_sxm/sparse_attention/vllm/0.24.0/dsv4_hca_attn_module_perf.parquet and /dev/null differ diff --git a/python/aisimulate/src/aisimulate_core/systems/data/b200_sxm/sparse_attention/vllm/0.25.0/collection_meta.yaml b/python/aisimulate/src/aisimulate_core/systems/data/b200_sxm/sparse_attention/vllm/0.25.0/collection_meta.yaml index 5b44fc818f..a9ed5fd7b3 100644 --- a/python/aisimulate/src/aisimulate_core/systems/data/b200_sxm/sparse_attention/vllm/0.25.0/collection_meta.yaml +++ b/python/aisimulate/src/aisimulate_core/systems/data/b200_sxm/sparse_attention/vllm/0.25.0/collection_meta.yaml @@ -40,22 +40,6 @@ tables: collected_at: '2026-09-15' rows: 768 status: complete - dsv4_hca_attn_module_perf: - rows: 3918 - status: complete - collections: - - collector_ref: c79cc132552deb677becc37f8139b0cca2b08d75 - collector_hash: sha256:ff1a5d288300da98a3c6f38c4894c17ad896f05384afcac3d966971f27130b83 - case_plan_hash: sha256:ff311865487afeebfbf526849e84ff3749335201ae74873e572766a06f91c63e - collected_at: '2026-09-15' - rows: 1960 - status: complete - - collector_ref: c79cc132552deb677becc37f8139b0cca2b08d75 - collector_hash: sha256:ff1a5d288300da98a3c6f38c4894c17ad896f05384afcac3d966971f27130b83 - case_plan_hash: sha256:e6845638359ae7fe7aa5ec93a5aa90bac79287cb3cb80368d0e4a7ac814ea456 - collected_at: '2026-09-15' - rows: 1958 - status: complete dsv4_hca_context_module_perf: rows: 5840 status: complete diff --git a/python/aisimulate/src/aisimulate_core/systems/data/b200_sxm/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet b/python/aisimulate/src/aisimulate_core/systems/data/b200_sxm/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet deleted file mode 100644 index 5057274a80..0000000000 Binary files a/python/aisimulate/src/aisimulate_core/systems/data/b200_sxm/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet and /dev/null differ diff --git a/python/aisimulate/src/aisimulate_core/systems/data/b200_sxm/vllm-0.25.0-collection-report.json b/python/aisimulate/src/aisimulate_core/systems/data/b200_sxm/vllm-0.25.0-collection-report.json index 1ed5848afb..4040b035ad 100644 --- a/python/aisimulate/src/aisimulate_core/systems/data/b200_sxm/vllm-0.25.0-collection-report.json +++ b/python/aisimulate/src/aisimulate_core/systems/data/b200_sxm/vllm-0.25.0-collection-report.json @@ -135,30 +135,6 @@ } ] }, - "dsv4_hca_attn_module_perf.parquet": { - "path": "sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet", - "rows": 3918, - "sha256": "dc93933dbfb7608991654d68500560a565b27976e4f10ff815d766f913566832", - "status": "complete", - "source_jobs": [ - { - "campaign": "campaign-ops-v1", - "shard": "dsv4_hca_attn_module-00", - "job": "job-1968531", - "collector_ref": "c79cc132552deb677becc37f8139b0cca2b08d75", - "source_rows": 1960, - "classified_failures": 0 - }, - { - "campaign": "campaign-ops-v1", - "shard": "dsv4_hca_attn_module-01", - "job": "job-1968532", - "collector_ref": "c79cc132552deb677becc37f8139b0cca2b08d75", - "source_rows": 1958, - "classified_failures": 0 - } - ] - }, "dsv4_hca_context_module_perf.parquet": { "path": "sparse_attention/vllm/0.25.0/dsv4_hca_context_module_perf.parquet", "rows": 5840, @@ -826,7 +802,8 @@ "computescale_perf.parquet": "Fresh smoke measured 5/5 cases but collector finalization rejected the existing undeclared scale_matrix_perf.txt side output. Original data and raw side output retained; do not claim end-to-end smoke passed.", "dsa_context_module_perf.parquet": "Full sweep incomplete; interrupted source snapshots have not finalized.", "dsa_generation_module_perf.parquet": "Full sweep incomplete; interrupted source snapshots have not finalized.", - "msa_generation_module_perf.parquet": "Full sweep incomplete; interrupted source snapshots have not finalized." + "msa_generation_module_perf.parquet": "Full sweep incomplete; interrupted source snapshots have not finalized.", + "dsv4_hca_attn_module_perf.parquet": "Published 2026-09-15 (3918 rows, jobs job-1968531/job-1968532, sha256 dc93933dbfb7608991654d68500560a565b27976e4f10ff815d766f913566832); removed 2026-10-15 together with the dsv4_hca_attn_module op: its SDK reader was a retired placeholder with no query path. The campaign ops/totals/clock entries above still record the run." }, "communication": "Raw CAR/NCCL results exist; not included in this operator-only publication batch.", "reused_kda": { @@ -836,8 +813,8 @@ "new_measurement": false, "sha256": "b21968ecb170a3095725a34b45cb16fe19c7dcd30250a04f5802d16a5d3c202a" }, - "total_published_rows": 419538, - "total_published_tables": 17, + "total_published_rows": 415620, + "total_published_tables": 16, "totals": { "planned": 426888, "successful": 342622, diff --git a/python/aisimulate/src/aisimulate_core/systems/data/b300_sxm/sparse_attention/vllm/0.22.0/reuse.yaml b/python/aisimulate/src/aisimulate_core/systems/data/b300_sxm/sparse_attention/vllm/0.22.0/reuse.yaml index 44542fcaab..abbd9ab01a 100644 --- a/python/aisimulate/src/aisimulate_core/systems/data/b300_sxm/sparse_attention/vllm/0.22.0/reuse.yaml +++ b/python/aisimulate/src/aisimulate_core/systems/data/b300_sxm/sparse_attention/vllm/0.22.0/reuse.yaml @@ -19,10 +19,6 @@ reuse: from_version: 0.24.0 reason: migrated from SHARED_LAYER_REUSE.txt (legacy newest-first effective donor) approved_by: yimingl -- table: dsv4_hca_attn_module_perf - from_version: 0.24.0 - reason: migrated from SHARED_LAYER_REUSE.txt (legacy newest-first effective donor) - approved_by: yimingl - table: dsv4_hca_context_module_perf from_version: 0.24.0 reason: migrated from SHARED_LAYER_REUSE.txt (legacy newest-first effective donor) diff --git a/python/aisimulate/src/aisimulate_core/systems/data/b300_sxm/sparse_attention/vllm/0.24.0/collection_meta.yaml b/python/aisimulate/src/aisimulate_core/systems/data/b300_sxm/sparse_attention/vllm/0.24.0/collection_meta.yaml index 8f27757676..71654ff7b6 100644 --- a/python/aisimulate/src/aisimulate_core/systems/data/b300_sxm/sparse_attention/vllm/0.24.0/collection_meta.yaml +++ b/python/aisimulate/src/aisimulate_core/systems/data/b300_sxm/sparse_attention/vllm/0.24.0/collection_meta.yaml @@ -27,8 +27,6 @@ tables: status: complete dsv4_csa_generation_module_perf: status: complete - dsv4_hca_attn_module_perf: - status: complete dsv4_hca_context_module_perf: status: complete dsv4_hca_generation_module_perf: diff --git a/python/aisimulate/src/aisimulate_core/systems/data/b300_sxm/sparse_attention/vllm/0.24.0/dsv4_hca_attn_module_perf.parquet b/python/aisimulate/src/aisimulate_core/systems/data/b300_sxm/sparse_attention/vllm/0.24.0/dsv4_hca_attn_module_perf.parquet deleted file mode 100644 index 40559a6601..0000000000 Binary files a/python/aisimulate/src/aisimulate_core/systems/data/b300_sxm/sparse_attention/vllm/0.24.0/dsv4_hca_attn_module_perf.parquet and /dev/null differ diff --git a/python/aisimulate/src/aisimulate_core/systems/data/b300_sxm/sparse_attention/vllm/0.25.0/collection_meta.yaml b/python/aisimulate/src/aisimulate_core/systems/data/b300_sxm/sparse_attention/vllm/0.25.0/collection_meta.yaml index 890fb99228..54b7e82d05 100644 --- a/python/aisimulate/src/aisimulate_core/systems/data/b300_sxm/sparse_attention/vllm/0.25.0/collection_meta.yaml +++ b/python/aisimulate/src/aisimulate_core/systems/data/b300_sxm/sparse_attention/vllm/0.25.0/collection_meta.yaml @@ -29,16 +29,6 @@ tables: collected_at: '2026-09-16' rows: 1544 status: complete - dsv4_hca_attn_module_perf: - rows: 3918 - status: complete - collections: - - collector_ref: cbaf51b64fa460e5ec6146bde407a4c64958212d - collector_hash: sha256:5882781b43c67c93f89e3e238978f461e242c83f681effd9ee6129a4be1637eb - case_plan_hash: sha256:0e844f577823301a0af954d0d20baab4642bb1c0368652998bc28160dfab816d - collected_at: '2026-09-16' - rows: 3918 - status: complete dsv4_hca_context_module_perf: rows: 5840 status: complete diff --git a/python/aisimulate/src/aisimulate_core/systems/data/b300_sxm/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet b/python/aisimulate/src/aisimulate_core/systems/data/b300_sxm/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet deleted file mode 100644 index 59d4ba25b4..0000000000 Binary files a/python/aisimulate/src/aisimulate_core/systems/data/b300_sxm/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet and /dev/null differ diff --git a/python/aisimulate/src/aisimulate_core/systems/data/gb200/sparse_attention/vllm/0.24.0/collection_meta.yaml b/python/aisimulate/src/aisimulate_core/systems/data/gb200/sparse_attention/vllm/0.24.0/collection_meta.yaml index 2bc2e20cb5..902301824b 100644 --- a/python/aisimulate/src/aisimulate_core/systems/data/gb200/sparse_attention/vllm/0.24.0/collection_meta.yaml +++ b/python/aisimulate/src/aisimulate_core/systems/data/gb200/sparse_attention/vllm/0.24.0/collection_meta.yaml @@ -27,8 +27,6 @@ tables: status: complete dsv4_csa_generation_module_perf: status: complete - dsv4_hca_attn_module_perf: - status: complete dsv4_hca_context_module_perf: status: complete dsv4_hca_generation_module_perf: diff --git a/python/aisimulate/src/aisimulate_core/systems/data/gb200/sparse_attention/vllm/0.24.0/dsv4_hca_attn_module_perf.parquet b/python/aisimulate/src/aisimulate_core/systems/data/gb200/sparse_attention/vllm/0.24.0/dsv4_hca_attn_module_perf.parquet deleted file mode 100644 index 1c1dd4d8e6..0000000000 Binary files a/python/aisimulate/src/aisimulate_core/systems/data/gb200/sparse_attention/vllm/0.24.0/dsv4_hca_attn_module_perf.parquet and /dev/null differ diff --git a/python/aisimulate/src/aisimulate_core/systems/data/gb200/sparse_attention/vllm/0.25.0/collection_meta.yaml b/python/aisimulate/src/aisimulate_core/systems/data/gb200/sparse_attention/vllm/0.25.0/collection_meta.yaml index df90df3d9c..e1618e5f12 100644 --- a/python/aisimulate/src/aisimulate_core/systems/data/gb200/sparse_attention/vllm/0.25.0/collection_meta.yaml +++ b/python/aisimulate/src/aisimulate_core/systems/data/gb200/sparse_attention/vllm/0.25.0/collection_meta.yaml @@ -29,16 +29,6 @@ tables: collected_at: '2026-09-17' rows: 1544 status: complete - dsv4_hca_attn_module_perf: - rows: 3918 - status: complete - collections: - - collector_ref: cbaf51b64fa460e5ec6146bde407a4c64958212d - collector_hash: sha256:5882781b43c67c93f89e3e238978f461e242c83f681effd9ee6129a4be1637eb - case_plan_hash: sha256:0e844f577823301a0af954d0d20baab4642bb1c0368652998bc28160dfab816d - collected_at: '2026-09-17' - rows: 3918 - status: complete dsv4_hca_context_module_perf: rows: 5840 status: complete diff --git a/python/aisimulate/src/aisimulate_core/systems/data/gb200/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet b/python/aisimulate/src/aisimulate_core/systems/data/gb200/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet deleted file mode 100644 index a57023a83d..0000000000 Binary files a/python/aisimulate/src/aisimulate_core/systems/data/gb200/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet and /dev/null differ diff --git a/python/aisimulate/src/aisimulate_core/systems/data/gb300/sparse_attention/vllm/0.22.0/reuse.yaml b/python/aisimulate/src/aisimulate_core/systems/data/gb300/sparse_attention/vllm/0.22.0/reuse.yaml index 44542fcaab..abbd9ab01a 100644 --- a/python/aisimulate/src/aisimulate_core/systems/data/gb300/sparse_attention/vllm/0.22.0/reuse.yaml +++ b/python/aisimulate/src/aisimulate_core/systems/data/gb300/sparse_attention/vllm/0.22.0/reuse.yaml @@ -19,10 +19,6 @@ reuse: from_version: 0.24.0 reason: migrated from SHARED_LAYER_REUSE.txt (legacy newest-first effective donor) approved_by: yimingl -- table: dsv4_hca_attn_module_perf - from_version: 0.24.0 - reason: migrated from SHARED_LAYER_REUSE.txt (legacy newest-first effective donor) - approved_by: yimingl - table: dsv4_hca_context_module_perf from_version: 0.24.0 reason: migrated from SHARED_LAYER_REUSE.txt (legacy newest-first effective donor) diff --git a/python/aisimulate/src/aisimulate_core/systems/data/gb300/sparse_attention/vllm/0.24.0/collection_meta.yaml b/python/aisimulate/src/aisimulate_core/systems/data/gb300/sparse_attention/vllm/0.24.0/collection_meta.yaml index 6369b50d82..eb86b23a6f 100644 --- a/python/aisimulate/src/aisimulate_core/systems/data/gb300/sparse_attention/vllm/0.24.0/collection_meta.yaml +++ b/python/aisimulate/src/aisimulate_core/systems/data/gb300/sparse_attention/vllm/0.24.0/collection_meta.yaml @@ -27,8 +27,6 @@ tables: status: complete dsv4_csa_generation_module_perf: status: complete - dsv4_hca_attn_module_perf: - status: complete dsv4_hca_context_module_perf: status: complete dsv4_hca_generation_module_perf: diff --git a/python/aisimulate/src/aisimulate_core/systems/data/gb300/sparse_attention/vllm/0.24.0/dsv4_hca_attn_module_perf.parquet b/python/aisimulate/src/aisimulate_core/systems/data/gb300/sparse_attention/vllm/0.24.0/dsv4_hca_attn_module_perf.parquet deleted file mode 100644 index 5b7342b5c7..0000000000 Binary files a/python/aisimulate/src/aisimulate_core/systems/data/gb300/sparse_attention/vllm/0.24.0/dsv4_hca_attn_module_perf.parquet and /dev/null differ diff --git a/python/aisimulate/src/aisimulate_core/systems/data/gb300/sparse_attention/vllm/0.25.0/collection_meta.yaml b/python/aisimulate/src/aisimulate_core/systems/data/gb300/sparse_attention/vllm/0.25.0/collection_meta.yaml index df90df3d9c..e1618e5f12 100644 --- a/python/aisimulate/src/aisimulate_core/systems/data/gb300/sparse_attention/vllm/0.25.0/collection_meta.yaml +++ b/python/aisimulate/src/aisimulate_core/systems/data/gb300/sparse_attention/vllm/0.25.0/collection_meta.yaml @@ -29,16 +29,6 @@ tables: collected_at: '2026-09-17' rows: 1544 status: complete - dsv4_hca_attn_module_perf: - rows: 3918 - status: complete - collections: - - collector_ref: cbaf51b64fa460e5ec6146bde407a4c64958212d - collector_hash: sha256:5882781b43c67c93f89e3e238978f461e242c83f681effd9ee6129a4be1637eb - case_plan_hash: sha256:0e844f577823301a0af954d0d20baab4642bb1c0368652998bc28160dfab816d - collected_at: '2026-09-17' - rows: 3918 - status: complete dsv4_hca_context_module_perf: rows: 5840 status: complete diff --git a/python/aisimulate/src/aisimulate_core/systems/data/gb300/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet b/python/aisimulate/src/aisimulate_core/systems/data/gb300/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet deleted file mode 100644 index f04505c645..0000000000 Binary files a/python/aisimulate/src/aisimulate_core/systems/data/gb300/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet and /dev/null differ diff --git a/python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.22.0/reuse.yaml b/python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.22.0/reuse.yaml index 44542fcaab..abbd9ab01a 100644 --- a/python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.22.0/reuse.yaml +++ b/python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.22.0/reuse.yaml @@ -19,10 +19,6 @@ reuse: from_version: 0.24.0 reason: migrated from SHARED_LAYER_REUSE.txt (legacy newest-first effective donor) approved_by: yimingl -- table: dsv4_hca_attn_module_perf - from_version: 0.24.0 - reason: migrated from SHARED_LAYER_REUSE.txt (legacy newest-first effective donor) - approved_by: yimingl - table: dsv4_hca_context_module_perf from_version: 0.24.0 reason: migrated from SHARED_LAYER_REUSE.txt (legacy newest-first effective donor) diff --git a/python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.24.0/collection_meta.yaml b/python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.24.0/collection_meta.yaml index 54bc3ce892..ea336b165f 100644 --- a/python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.24.0/collection_meta.yaml +++ b/python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.24.0/collection_meta.yaml @@ -27,8 +27,6 @@ tables: status: complete dsv4_csa_generation_module_perf: status: complete - dsv4_hca_attn_module_perf: - status: complete dsv4_hca_context_module_perf: status: complete dsv4_hca_generation_module_perf: diff --git a/python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.24.0/dsv4_hca_attn_module_perf.parquet b/python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.24.0/dsv4_hca_attn_module_perf.parquet deleted file mode 100644 index dd3336ef7c..0000000000 Binary files a/python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.24.0/dsv4_hca_attn_module_perf.parquet and /dev/null differ diff --git a/python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.25.0/collection_meta.yaml b/python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.25.0/collection_meta.yaml index acb9dd93ed..20b19499db 100644 --- a/python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.25.0/collection_meta.yaml +++ b/python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.25.0/collection_meta.yaml @@ -35,16 +35,6 @@ tables: collected_at: '2026-09-17' rows: 1544 status: complete - dsv4_hca_attn_module_perf: - rows: 3918 - status: complete - collections: - - collector_ref: cbaf51b64fa460e5ec6146bde407a4c64958212d - collector_hash: sha256:5882781b43c67c93f89e3e238978f461e242c83f681effd9ee6129a4be1637eb - case_plan_hash: sha256:0e844f577823301a0af954d0d20baab4642bb1c0368652998bc28160dfab816d - collected_at: '2026-09-17' - rows: 3918 - status: complete dsv4_hca_context_module_perf: rows: 5793 status: complete diff --git a/python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet b/python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet deleted file mode 100644 index b105008d6a..0000000000 Binary files a/python/aisimulate/src/aisimulate_core/systems/data/h100_sxm/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet and /dev/null differ diff --git a/python/aisimulate/src/aisimulate_core/systems/data/h200_sxm/sparse_attention/vllm/0.24.0/collection_meta.yaml b/python/aisimulate/src/aisimulate_core/systems/data/h200_sxm/sparse_attention/vllm/0.24.0/collection_meta.yaml index 0a65ebbc05..25ecac32f0 100644 --- a/python/aisimulate/src/aisimulate_core/systems/data/h200_sxm/sparse_attention/vllm/0.24.0/collection_meta.yaml +++ b/python/aisimulate/src/aisimulate_core/systems/data/h200_sxm/sparse_attention/vllm/0.24.0/collection_meta.yaml @@ -27,8 +27,6 @@ tables: status: complete dsv4_csa_generation_module_perf: status: complete - dsv4_hca_attn_module_perf: - status: complete dsv4_hca_context_module_perf: status: complete dsv4_hca_generation_module_perf: diff --git a/python/aisimulate/src/aisimulate_core/systems/data/h200_sxm/sparse_attention/vllm/0.24.0/dsv4_hca_attn_module_perf.parquet b/python/aisimulate/src/aisimulate_core/systems/data/h200_sxm/sparse_attention/vllm/0.24.0/dsv4_hca_attn_module_perf.parquet deleted file mode 100644 index a33c64b8f5..0000000000 Binary files a/python/aisimulate/src/aisimulate_core/systems/data/h200_sxm/sparse_attention/vllm/0.24.0/dsv4_hca_attn_module_perf.parquet and /dev/null differ diff --git a/python/aisimulate/src/aisimulate_core/systems/data/h200_sxm/sparse_attention/vllm/0.25.0/collection_meta.yaml b/python/aisimulate/src/aisimulate_core/systems/data/h200_sxm/sparse_attention/vllm/0.25.0/collection_meta.yaml index 0189e099c3..b11bfba80e 100644 --- a/python/aisimulate/src/aisimulate_core/systems/data/h200_sxm/sparse_attention/vllm/0.25.0/collection_meta.yaml +++ b/python/aisimulate/src/aisimulate_core/systems/data/h200_sxm/sparse_attention/vllm/0.25.0/collection_meta.yaml @@ -29,16 +29,6 @@ tables: collected_at: '2026-09-17' rows: 1544 status: complete - dsv4_hca_attn_module_perf: - rows: 3918 - status: complete - collections: - - collector_ref: cbaf51b64fa460e5ec6146bde407a4c64958212d - collector_hash: sha256:5882781b43c67c93f89e3e238978f461e242c83f681effd9ee6129a4be1637eb - case_plan_hash: sha256:0e844f577823301a0af954d0d20baab4642bb1c0368652998bc28160dfab816d - collected_at: '2026-09-17' - rows: 3918 - status: complete dsv4_hca_context_module_perf: rows: 5840 status: complete diff --git a/python/aisimulate/src/aisimulate_core/systems/data/h200_sxm/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet b/python/aisimulate/src/aisimulate_core/systems/data/h200_sxm/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet deleted file mode 100644 index d865bf5b4c..0000000000 Binary files a/python/aisimulate/src/aisimulate_core/systems/data/h200_sxm/sparse_attention/vllm/0.25.0/dsv4_hca_attn_module_perf.parquet and /dev/null differ diff --git a/python/aisimulate/src/aisimulate_core/systems/perf_data_reuse_manifest.yaml b/python/aisimulate/src/aisimulate_core/systems/perf_data_reuse_manifest.yaml index d7fa47380d..446e4f4d45 100644 --- a/python/aisimulate/src/aisimulate_core/systems/perf_data_reuse_manifest.yaml +++ b/python/aisimulate/src/aisimulate_core/systems/perf_data_reuse_manifest.yaml @@ -414,11 +414,6 @@ groups: tier: shared systems: [b200_sxm, b300_sxm, gb200, gb300, h100_sxm, h200_sxm] frameworks: [sglang] - - op_file: dsv4_hca_attn_module_perf.parquet - kernel_source: 'FLASHMLA_SPARSE_DSV4' - tier: shared - systems: [b200_sxm, b300_sxm, gb200, gb300, h100_sxm, h200_sxm] - frameworks: [vllm] - op_file: dsv4_hca_context_module_perf.parquet kernel_source: 'DeepseekV4TrtllmAttention' tier: shared diff --git a/python/aisimulate/tests/golden/generator/slurm/sglang/0.5.11/agg.yaml b/python/aisimulate/tests/golden/generator/slurm/sglang/0.5.11/agg.yaml index 94126be193..de4bf7ad99 100644 --- a/python/aisimulate/tests/golden/generator/slurm/sglang/0.5.11/agg.yaml +++ b/python/aisimulate/tests/golden/generator/slurm/sglang/0.5.11/agg.yaml @@ -35,7 +35,7 @@ deployment.json: |- "--max-running-requests", "512", "--max-prefill-tokens", - "1628", + "1664", "--enable-mixed-chunk", "--cuda-graph-bs", "1", diff --git a/python/aisimulate/tests/golden/generator/slurm/sglang/0.5.11/disagg.yaml b/python/aisimulate/tests/golden/generator/slurm/sglang/0.5.11/disagg.yaml index d5e774b21b..17e5ee5747 100644 --- a/python/aisimulate/tests/golden/generator/slurm/sglang/0.5.11/disagg.yaml +++ b/python/aisimulate/tests/golden/generator/slurm/sglang/0.5.11/disagg.yaml @@ -35,7 +35,7 @@ deployment.json: |- "--max-running-requests", "8", "--max-prefill-tokens", - "1628", + "1664", "--disaggregation-transfer-backend", "nixl", "--disaggregation-mode", @@ -72,7 +72,7 @@ deployment.json: |- "--max-running-requests", "512", "--max-prefill-tokens", - "1628", + "1664", "--disaggregation-transfer-backend", "nixl", "--cuda-graph-bs", diff --git a/python/aisimulate/tests/golden/generator/slurm/trtllm/1.3.0rc14/disagg.yaml b/python/aisimulate/tests/golden/generator/slurm/trtllm/1.3.0rc14/disagg.yaml index 03ac01d2c4..21973c615f 100644 --- a/python/aisimulate/tests/golden/generator/slurm/trtllm/1.3.0rc14/disagg.yaml +++ b/python/aisimulate/tests/golden/generator/slurm/trtllm/1.3.0rc14/disagg.yaml @@ -65,7 +65,7 @@ deployment.json: |- "--max-batch-size", "8", "--max-num-tokens", - "32", + "64", "--extra-engine-args", "decode_config.yaml", "--gpus-per-node", @@ -129,7 +129,7 @@ decode_config.yaml: | enable_attention_dp: false enable_chunked_prefill: false max_batch_size: 8 - max_num_tokens: 32 + max_num_tokens: 64 kv_cache_config: free_gpu_memory_fraction: 0.8 dtype: auto diff --git a/python/aisimulate/tests/golden/generator/slurm/vllm/0.20.1/agg.yaml b/python/aisimulate/tests/golden/generator/slurm/vllm/0.20.1/agg.yaml index 58e9faf3f0..689c37d6cc 100644 --- a/python/aisimulate/tests/golden/generator/slurm/vllm/0.20.1/agg.yaml +++ b/python/aisimulate/tests/golden/generator/slurm/vllm/0.20.1/agg.yaml @@ -37,7 +37,7 @@ deployment.json: |- "--max-num-seqs", "512", "--max-num-batched-tokens", - "2140" + "2176" ], "env": { "VLLM_DP_SIZE": "1", diff --git a/python/aisimulate/tests/golden/generator/slurm/vllm/0.20.1/disagg.yaml b/python/aisimulate/tests/golden/generator/slurm/vllm/0.20.1/disagg.yaml index 57aaf7a702..f4d1f9cb95 100644 --- a/python/aisimulate/tests/golden/generator/slurm/vllm/0.20.1/disagg.yaml +++ b/python/aisimulate/tests/golden/generator/slurm/vllm/0.20.1/disagg.yaml @@ -35,7 +35,7 @@ deployment.json: |- "--max-num-seqs", "8", "--max-num-batched-tokens", - "1628", + "1664", "--disaggregation-mode", "prefill", "--kv-transfer-config", diff --git a/python/aisimulate/tests/golden/generator/test_release_1_2_0_backend_contracts.py b/python/aisimulate/tests/golden/generator/test_release_1_2_0_backend_contracts.py index 0ef122af43..abebc367c6 100644 --- a/python/aisimulate/tests/golden/generator/test_release_1_2_0_backend_contracts.py +++ b/python/aisimulate/tests/golden/generator/test_release_1_2_0_backend_contracts.py @@ -313,7 +313,8 @@ def test_vllm_0_20_1_cli_args_golden_contract(): assert _value_after(tokens, "--data-parallel-size") == "1" assert _value_after(tokens, "--kv-cache-dtype") == "bfloat16" assert _value_after(tokens, "--gpu-memory-utilization") == "0.82" - assert _value_after(tokens, "--max-num-batched-tokens") == "4060" + # token budgets round up to a multiple of 64 (backend_config_mapping.yaml): 4060 -> 4096 + assert _value_after(tokens, "--max-num-batched-tokens") == "4096" assert "--enable-expert-parallel" in flags assert "--enforce-eager" in flags assert "--no-enable-prefix-caching" not in flags @@ -362,7 +363,8 @@ def test_sglang_0_5_11_cli_args_golden_contract(): assert _value_after(tokens, "--kv-cache-dtype") == "auto" assert _value_after(tokens, "--mem-fraction-static") == "0.82" assert _value_after(tokens, "--chunked-prefill-size") == "-1" - assert _value_after(tokens, "--max-prefill-tokens") == "3548" + # token budgets round up to a multiple of 64 (backend_config_mapping.yaml): 3548 -> 3584 + assert _value_after(tokens, "--max-prefill-tokens") == "3584" assert _value_after(tokens, "--context-length") == "4096" assert _value_after(tokens, "--max-running-requests") == "512" assert _value_after(tokens, "--speculative-algorithm") == "NEXTN" diff --git a/python/aisimulate/tests/unit/collector/sglang/test_collect_attn_encoder_sm_dispatch.py b/python/aisimulate/tests/unit/collector/sglang/test_collect_attn_encoder_sm_dispatch.py new file mode 100644 index 0000000000..df6d1eb3e5 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/sglang/test_collect_attn_encoder_sm_dispatch.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The encoder collector picks the vision-attention lane the way sglang does: by capability MAJOR +(vision.py:1298-1303 @0.5.21), so SM103 (B300) shares FA4 with SM100 instead of falling to Triton.""" + +import ast +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.unit +SOURCE = Path(__file__).resolve().parents[4] / "collector" / "sglang" / "collect_attn_encoder.py" + + +def _serving_vision_backend(): + tree = ast.parse(SOURCE.read_text(encoding="utf-8")) + fn = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "serving_vision_backend") + ns: dict = {} + exec(compile(ast.Module(body=[fn], type_ignores=[]), str(SOURCE), "exec"), ns) + return ns["serving_vision_backend"] + + +@pytest.mark.parametrize( + "sm, expected", + [(90, "fa3"), (100, "fa4"), (103, "fa4"), (120, "triton_attn")], +) +def test_vision_backend_follows_capability_major(sm, expected): + assert _serving_vision_backend()(sm) == expected + + +def test_run_dispatch_uses_the_family_helper_not_exact_sm(): + src = SOURCE.read_text(encoding="utf-8") + assert "if sm == 100" not in src and "if sm == 90" not in src + assert "backend = serving_vision_backend(sm)" in src diff --git a/python/aisimulate/tests/unit/collector/sglang/test_collect_attn_flashinfer_workspace.py b/python/aisimulate/tests/unit/collector/sglang/test_collect_attn_flashinfer_workspace.py new file mode 100644 index 0000000000..5acf02b7a3 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/sglang/test_collect_attn_flashinfer_workspace.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The FlashInfer workspace floor goes through the knob sglang 0.5.21 reads +(envs.SGLANG_FLASHINFER_WORKSPACE_SIZE), not the retired global_config attribute +(ed9109ec wrote the latter and changed nothing: l40s pipeline 72343559).""" + +import ast +import sys +import types +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.unit +SOURCE = Path(__file__).resolve().parents[4] / "collector" / "sglang" / "collect_attn.py" + + +def _load(): + tree = ast.parse(SOURCE.read_text(encoding="utf-8"), filename=str(SOURCE)) + nodes = [ + n + for n in tree.body + if (isinstance(n, ast.FunctionDef) and n.name == "_ensure_flashinfer_workspace") + or ( + isinstance(n, ast.Assign) + and any(getattr(t, "id", "") == "_FLASHINFER_WORKSPACE_MIN_BYTES" for t in n.targets) + ) + ] + ns = {} + exec(compile(ast.Module(body=nodes, type_ignores=[]), str(SOURCE), "exec"), ns) + return ns["_ensure_flashinfer_workspace"], ns["_FLASHINFER_WORKSPACE_MIN_BYTES"] + + +class _EnvInt: + def __init__(self, v): + self.v = v + + def get(self): + return self.v + + def set(self, v): + self.v = v + + +@pytest.mark.parametrize(("initial", "expected"), [(384 * 1024 * 1024, 4 << 30), (8 << 30, 8 << 30)]) +def test_0521_envs_knob_is_raised_to_the_floor_and_a_larger_setting_is_kept(monkeypatch, initial, expected): + fn, floor = _load() + assert floor == 4 << 30 + knob = _EnvInt(initial) + monkeypatch.setitem(sys.modules, "sglang", types.ModuleType("sglang")) + monkeypatch.setitem(sys.modules, "sglang.srt", types.ModuleType("sglang.srt")) + environ = types.ModuleType("sglang.srt.environ") + environ.envs = types.SimpleNamespace(SGLANG_FLASHINFER_WORKSPACE_SIZE=knob) + monkeypatch.setitem(sys.modules, "sglang.srt.environ", environ) + assert fn() == expected and knob.get() == expected diff --git a/python/aisimulate/tests/unit/collector/sglang/test_collect_gdn_contract.py b/python/aisimulate/tests/unit/collector/sglang/test_collect_gdn_contract.py index 8403ec4386..c3634dbed9 100644 --- a/python/aisimulate/tests/unit/collector/sglang/test_collect_gdn_contract.py +++ b/python/aisimulate/tests/unit/collector/sglang/test_collect_gdn_contract.py @@ -324,6 +324,7 @@ def recording_zeros(*args, **kwargs): "fused_recurrent_gated_delta_rule_packed_decode" ), "_resolve_flashinfer_gdn_decode": resolve or (lambda _dtype: (flashinfer_kernel, flashinfer_error)), + "_resolve_fused_gdn_decode_conv": lambda: None, # <0.5.21 split conv lane "benchmark_with_power": lambda *, kernel_func, **_kwargs: _FakeBenchmark(kernel_func), "log_perf": lambda **kwargs: logged.append(kwargs["kernel_source"]) or True, }, diff --git a/python/aisimulate/tests/unit/collector/sglang/test_collect_kda_contract.py b/python/aisimulate/tests/unit/collector/sglang/test_collect_kda_contract.py index e867243d64..26376a8405 100644 --- a/python/aisimulate/tests/unit/collector/sglang/test_collect_kda_contract.py +++ b/python/aisimulate/tests/unit/collector/sglang/test_collect_kda_contract.py @@ -82,21 +82,6 @@ def test_glm5_next_decode_uses_bounded_gate_recurrence_not_packed_decode(): assert "fused_recurrent_kda_packed_decode" not in names -@pytest.mark.parametrize( - ("version", "accepted"), - [("0.5.16", True), ("0.5.17", False), ("0.5.19", False), ("0.5.20", True), ("0.5.21", False)], -) -def test_kda_compat_admits_kimi_branch_and_glm_runtime(version, accepted): - from collector.version_resolver import _check_compat - - declaration = next( - node.value.value - for node in ast.parse(SOURCE_PATH.read_text(encoding="utf-8")).body - if isinstance(node, ast.Assign) and any(getattr(t, "id", None) == "__compat__" for t in node.targets) - ) - assert _check_compat(declaration, version) is accepted - - def test_glm5_next_chunk_kda_int32_offset_guard_precedes_the_scan(): # chunk_kda's Triton kernels use int32 (bos * H + i_h) * K offsets # (fla/kda.py:265-268 @v0.5.20); an overflowing launch is an illegal @@ -134,23 +119,14 @@ def _gate_namespace(): return namespace -def test_kda_audited_release_sets_do_not_overlap(): - from collector.version_resolver import _check_compat - - compat = _gate_namespace()["_KDA_ARCHITECTURE_COMPAT"] - audited = {arch: spec.split("==", 1)[1] for arch, spec in compat.items()} - assert set(audited) == {"KimiK3ForConditionalGeneration", "Glm5NextForConditionalGeneration"} - for arch, version in audited.items(): - for other, spec in compat.items(): - assert _check_compat(spec, version) is (other == arch), (arch, other) - - @pytest.mark.parametrize( ("model_name", "version", "audited"), [ - ("moonshotai/Kimi-K3", "0.5.16", True), + ("moonshotai/Kimi-K3", "0.5.16", False), + ("moonshotai/Kimi-K3", "0.5.21", True), ("moonshotai/Kimi-K3", "0.5.20", False), - ("zai-org/GLM-5.3-Flash", "0.5.20", True), + ("zai-org/GLM-5.3-Flash", "0.5.21", False), + ("zai-org/GLM-5.3-Flash", "0.5.20", False), # 0.5.20 lane retired 2026-10-30 ("zai-org/GLM-5.3-Flash", "0.5.16", False), ], ) diff --git a/python/aisimulate/tests/unit/collector/sglang/test_collect_mhc_module.py b/python/aisimulate/tests/unit/collector/sglang/test_collect_mhc_module.py index 9ae8dc5b71..65a463f094 100644 --- a/python/aisimulate/tests/unit/collector/sglang/test_collect_mhc_module.py +++ b/python/aisimulate/tests/unit/collector/sglang/test_collect_mhc_module.py @@ -50,7 +50,7 @@ def test_mhc_sweep_fails_closed_after_an_inner_error(monkeypatch): runner = SimpleNamespace(model=SimpleNamespace(model=SimpleNamespace(layers=[layer]))) monkeypatch.setattr(module, "_load_one_layer_runner", lambda *_args, **_kwargs: runner) monkeypatch.setattr(module, "_hidden_size", lambda _layer: 4096) - monkeypatch.setattr(module, "get_version", lambda _package: "0.5.14") + monkeypatch.setattr(module, "get_version", lambda _package: "0.5.21") monkeypatch.setattr(module, "_make_residual", lambda _layer, num_tokens, _device: num_tokens) monkeypatch.setattr(module, "_make_kernel", lambda _layer, _op, residual: residual) monkeypatch.setattr(module, "_mhc_call_args", lambda _layer: (None, None)) diff --git a/python/aisimulate/tests/unit/collector/sglang/test_collect_mla_graph_capture.py b/python/aisimulate/tests/unit/collector/sglang/test_collect_mla_graph_capture.py new file mode 100644 index 0000000000..a68041f4cf --- /dev/null +++ b/python/aisimulate/tests/unit/collector/sglang/test_collect_mla_graph_capture.py @@ -0,0 +1,175 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""sglang mla_context on SM100 is timed under CUDA-graph replay again. + +Between 2026-10-05 and 2026-10-11 the trtllm_mla PREFILL lane was timed eagerly +because flashinfer 0.6.18 refuses to capture trtllm_ragged_attention_deepseek +without the CPU length mirrors sglang 0.5.21 does not pass; the eager rows +floored at the Python launch loop (B200 0.13 ms, GB200 0.29 ms for every +isl <= 1024 cell). The collector now hands flashinfer the mirrors from +sglang's own extend_seq_lens_cpu and refuses to publish any eager result.""" + +import ast +import contextlib +import inspect +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +pytestmark = pytest.mark.unit + +REPO_ROOT = Path(__file__).resolve().parents[4] +SOURCE = REPO_ROOT / "collector" / "sglang" / "collect_mla.py" +FORWARD_MODE = SimpleNamespace(EXTEND="extend", DECODE="decode") + + +def _tree(): + return ast.parse(SOURCE.read_text(encoding="utf-8"), filename=str(SOURCE)) + + +def _load(names, namespace): + selected = [ + node for node in _tree().body if isinstance(node, ast.FunctionDef | ast.ClassDef) and node.name in names + ] + assert len(selected) == len(names) + loaded = dict(namespace) + exec(compile(ast.Module(body=selected, type_ignores=[]), str(SOURCE), "exec"), loaded) + return loaded + + +def _fake_torch(): + return SimpleNamespace( + int32="int32", + tensor=lambda data, dtype=None, device=None: SimpleNamespace(values=tuple(data), dtype=dtype, device=device), + ) + + +def _install_fake_flashinfer(monkeypatch): + calls = [] + + def trtllm_ragged_attention_deepseek( + *, query, cum_seq_lens_q, cum_seq_lens_kv, q_seq_lens_cpu=None, kv_seq_lens_cpu=None + ): + calls.append({"q_seq_lens_cpu": q_seq_lens_cpu, "kv_seq_lens_cpu": kv_seq_lens_cpu}) + return "out" + + prefill = SimpleNamespace(trtllm_ragged_attention_deepseek=trtllm_ragged_attention_deepseek) + monkeypatch.setitem(sys.modules, "flashinfer", SimpleNamespace(prefill=prefill)) + monkeypatch.setitem(sys.modules, "flashinfer.prefill", prefill) + return prefill, calls, trtllm_ragged_attention_deepseek + + +class _ServingBackend: + """Stand-in for sglang 0.5.21's TRTLLMMLABackend: metadata init + the kernel hook.""" + + def __init__(self): + self.metadata_batches = [] + + def init_forward_metadata(self, forward_batch): + self.metadata_batches.append(forward_batch) + + def _run_prefill_kernel(self, *, q, cum_seq_lens_q, cum_seq_lens_kv, **_rest): + # sglang calls flashinfer through the package attribute at call time + # (trtllm_mla_backend.py:1121 @v0.5.21), without the CPU mirrors. + import flashinfer + + return flashinfer.prefill.trtllm_ragged_attention_deepseek( + query=q, cum_seq_lens_q=cum_seq_lens_q, cum_seq_lens_kv=cum_seq_lens_kv + ) + + +def _backend(): + loaded = _load( + ["_extend_seq_lens_cpu", "_ragged_prefill_with_cpu_seq_lens", "_GraphCapturableTRTLLMMLABackend"], + { + "torch": _fake_torch(), + "ForwardMode": FORWARD_MODE, + "TRTLLMMLABackend": _ServingBackend, + "contextlib": contextlib, + "inspect": inspect, + }, + ) + return loaded["_GraphCapturableTRTLLMMLABackend"]() + + +def _extend_batch(lens): + return SimpleNamespace(forward_mode=FORWARD_MODE.EXTEND, extend_seq_lens_cpu=list(lens)) + + +def test_no_prefix_prefill_call_carries_the_cpu_length_mirrors(monkeypatch): + prefill, calls, original = _install_fake_flashinfer(monkeypatch) + backend = _backend() + backend.init_forward_metadata(_extend_batch([512, 512])) + cum = object() + + assert backend._run_prefill_kernel(q="q", cum_seq_lens_q=cum, cum_seq_lens_kv=cum) == "out" + + (call,) = calls + assert call["q_seq_lens_cpu"].values == (512, 512) + assert call["q_seq_lens_cpu"].dtype == "int32" + assert call["q_seq_lens_cpu"].device == "cpu" + assert call["kv_seq_lens_cpu"] is call["q_seq_lens_cpu"] + # the patch is scoped to the call: serving's function is back afterwards + assert prefill.trtllm_ragged_attention_deepseek is original + + +def test_chunked_prefix_call_goes_through_sglang_unchanged(monkeypatch): + _, calls, _ = _install_fake_flashinfer(monkeypatch) + backend = _backend() + backend.init_forward_metadata(_extend_batch([512])) + + backend._run_prefill_kernel(q="q", cum_seq_lens_q=object(), cum_seq_lens_kv=object()) + + assert calls == [{"q_seq_lens_cpu": None, "kv_seq_lens_cpu": None}] + + +def test_decode_metadata_leaves_the_prefill_hook_inert(monkeypatch): + _, calls, _ = _install_fake_flashinfer(monkeypatch) + backend = _backend() + backend.init_forward_metadata(SimpleNamespace(forward_mode=FORWARD_MODE.DECODE, seq_lens_cpu=[9])) + cum = object() + + backend._run_prefill_kernel(q="q", cum_seq_lens_q=cum, cum_seq_lens_kv=cum) + + assert calls == [{"q_seq_lens_cpu": None, "kv_seq_lens_cpu": None}] + + +def test_benchmark_layer_refuses_eager_timing(): + benchmark_kwargs = {} + + @contextlib.contextmanager + def benchmark_with_power(**kwargs): + benchmark_kwargs.update(kwargs) + yield {"latency_ms": 0.13, "power_stats": None, "used_cuda_graph": False} + + benchmark_layer = _load( + ["benchmark_layer"], + { + "benchmark_with_power": benchmark_with_power, + "forward_context": lambda _ctx: contextlib.nullcontext(), + "ForwardContext": lambda **_kw: None, + }, + )["benchmark_layer"] + forward_batch = SimpleNamespace(attn_backend=None) + + with pytest.raises(RuntimeError, match="refusing to publish eager timing"): + benchmark_layer(lambda *a, **k: None, forward_batch, SimpleNamespace(device="cuda:0"), None, None, None, None) + + # graph timing is the helper default; the lane must not opt out of it + assert "use_cuda_graph" not in benchmark_kwargs + assert benchmark_kwargs.get("allow_graph_fail") is not True + + +def test_run_mla_has_no_eager_prefill_switch(): + source = SOURCE.read_text(encoding="utf-8") + assert "use_cuda_graph=False" not in source + run_mla = next(node for node in _tree().body if isinstance(node, ast.FunctionDef) and node.name == "run_mla") + constructed = { + ast.unparse(node.func) + for node in ast.walk(run_mla) + if isinstance(node, ast.Call) and ast.unparse(node.func).endswith("TRTLLMMLABackend") + } + assert constructed == {"_GraphCapturableTRTLLMMLABackend"} diff --git a/python/aisimulate/tests/unit/collector/sglang/test_collect_mla_module.py b/python/aisimulate/tests/unit/collector/sglang/test_collect_mla_module.py index 6609218b6f..6663fa9dd8 100644 --- a/python/aisimulate/tests/unit/collector/sglang/test_collect_mla_module.py +++ b/python/aisimulate/tests/unit/collector/sglang/test_collect_mla_module.py @@ -9,6 +9,7 @@ from unittest.mock import patch import pytest +from collector.helper import PhaseClock, render_phase_totals, run_measurement_subprocess @pytest.fixture(autouse=True) @@ -19,6 +20,9 @@ def _mock_helper_imports(monkeypatch): fake_helper.get_sm_version = lambda: 90 # default Hopper fake_helper.log_perf = lambda **kw: None fake_helper.benchmark_with_power = lambda **kw: None + fake_helper.PhaseClock = PhaseClock + fake_helper.render_phase_totals = render_phase_totals + fake_helper.run_measurement_subprocess = run_measurement_subprocess monkeypatch.setitem(__import__("sys").modules, "helper", fake_helper) fake_torch = types.ModuleType("torch") @@ -149,10 +153,11 @@ def test_memory_guard(self): class TestDsaContextPrefixShape: - def test_accepts_single_token_extend_without_reclassifying_it_as_decode(self): + def test_single_token_extend_stays_out_while_prefill_is_timed_eagerly(self): + # 0.5.21: prefill graphs off, and the eager DSA indexer faulted on a one-token extend mod = _import_module() - assert mod._dsa_context_prefix_shape_is_valid(1, 1, 0) - assert mod._dsa_context_prefix_shape_is_valid(32, 1, 1024) + assert not mod._dsa_context_prefix_shape_is_valid(1, 1, 0) + assert not mod._dsa_context_prefix_shape_is_valid(32, 1, 1024) def test_accepts_multi_token_prefill_with_prefix(self): mod = _import_module() @@ -695,28 +700,16 @@ def test_zero_logged_rows_raises(self, monkeypatch): attention_backend="fa3", ) - def test_configured_subprocess_timeout_raises(self, monkeypatch): + def test_configured_subprocess_timeout_raises(self, monkeypatch, capsys): mod = _import_module() - class TimedOutProcess: - returncode = None - killed = False - waited = False + def timed_out(argv, *, env, cwd, timeout): + # helper.run_measurement_subprocess kills the child and keeps its output + assert timeout == 7 + raise mod.subprocess.TimeoutExpired(cmd=argv, timeout=timeout, output=b"partial sweep") - def communicate(self, timeout=None): - assert timeout == 7 - raise mod.subprocess.TimeoutExpired(cmd="mla-collector", timeout=timeout) - - def kill(self): - self.killed = True - - def wait(self): - self.waited = True - self.returncode = -9 - - process = TimedOutProcess() monkeypatch.setenv("AIC_MLA_MODULE_SUBPROCESS_TIMEOUT_SEC", "7") - monkeypatch.setattr(mod.subprocess, "Popen", lambda *_args, **_kwargs: process) + monkeypatch.setattr(mod, "run_measurement_subprocess", timed_out) with pytest.raises(RuntimeError, match=r"MLA module context subprocess timed out after 7s"): mod._run_mla_subprocess( @@ -731,8 +724,7 @@ def wait(self): attention_backend="fa3", ) - assert process.killed - assert process.waited + assert "partial sweep" in capsys.readouterr().out class TestOrdinaryMLAPrecision: @@ -772,14 +764,14 @@ def test_worker_chunk_size_reaches_subprocess_call(self, monkeypatch, chunk_size def capture_call(*_args, **kwargs): captured.update(kwargs) - def popen(argv, **_kwargs): + def run_child(argv, **_kwargs): # Execute the generated child call without importing the GPU runtime. tree = ast.parse(argv[2]) call = ast.Expression(tree.body[-1].value) eval(compile(call, "", "eval"), {"run_mla_module": capture_call}) - return types.SimpleNamespace(returncode=0, communicate=lambda timeout: (b"", None)) + return 0, b"" - monkeypatch.setattr(mod.subprocess, "Popen", popen) + monkeypatch.setattr(mod, "run_measurement_subprocess", run_child) mod.run_mla_module_worker( 128, 1, diff --git a/python/aisimulate/tests/unit/collector/sglang/test_collect_mla_module_native_quant.py b/python/aisimulate/tests/unit/collector/sglang/test_collect_mla_module_native_quant.py index 62b3911543..db5b095ad5 100644 --- a/python/aisimulate/tests/unit/collector/sglang/test_collect_mla_module_native_quant.py +++ b/python/aisimulate/tests/unit/collector/sglang/test_collect_mla_module_native_quant.py @@ -15,6 +15,7 @@ from unittest.mock import patch import pytest +from collector.helper import PhaseClock, render_phase_totals, run_measurement_subprocess pytestmark = pytest.mark.unit @@ -43,6 +44,10 @@ def _mock_helper_imports(monkeypatch): fake_helper.get_sm_version = lambda: 90 # default Hopper fake_helper.log_perf = lambda **kw: None fake_helper.benchmark_with_power = lambda **kw: None + # the collector imports the phase-timing helpers at module load + fake_helper.PhaseClock = PhaseClock + fake_helper.render_phase_totals = render_phase_totals + fake_helper.run_measurement_subprocess = run_measurement_subprocess monkeypatch.setitem(sys.modules, "helper", fake_helper) fake_torch = types.ModuleType("torch") diff --git a/python/aisimulate/tests/unit/collector/sglang/test_collect_moe_population.py b/python/aisimulate/tests/unit/collector/sglang/test_collect_moe_population.py index 2e9300a98f..9ca556fc29 100644 --- a/python/aisimulate/tests/unit/collector/sglang/test_collect_moe_population.py +++ b/python/aisimulate/tests/unit/collector/sglang/test_collect_moe_population.py @@ -167,9 +167,12 @@ def test_kimi_int4_population_restores_all_hopper_ep_slices(): module_config={"group_size": 32}, ) - assert len(populated) == 3078 - assert sum(case[7] == 64 for case in populated) == 243 - assert sum(case[7] == 128 for case in populated) == 162 + # 114 (quant, shape, tp, ep, distribution) series x 27 token points = the 3078 + # per-token rows the per-token layout used to schedule as separate cases. + assert len(populated) == 114 + assert sum(len(case[1]) for case in populated) == 3078 + assert sum(len(case[1]) for case in populated if case[7] == 64) == 243 + assert sum(len(case[1]) for case in populated if case[7] == 128) == 162 # int4_wo opened to SM100/103 (owner decision 2026-07-11): serving # auto-selects flashinfer_trtllm for Kimi INT4 there (server_args.py:3736), # replacing the old zero-population contract. @@ -179,7 +182,7 @@ def test_kimi_int4_population_restores_all_hopper_ep_slices(): allowed_mode="int4_wo", module_config={"group_size": 32}, ) - assert len(sm100) == 3078 + assert len(sm100) == 114 and sum(len(case[1]) for case in sm100) == 3078 from types import SimpleNamespace from collector.case_generator import get_sglang_moe_backend @@ -291,7 +294,7 @@ def test_dsv4_fp4_population_keeps_runtime_limits_attempted(): allowed_mode="w4a8_mxfp4_mxfp8", ) assert len(attempted) == 1 - assert attempted[0][1:8] == [16384, 7168, 3072, 6, 384, 16, 2] + assert attempted[0][1:8] == [[16384], 7168, 3072, 6, 384, 16, 2] def test_sm120_shared_memory_predictions_do_not_remove_moe_cases(): @@ -504,7 +507,9 @@ def test_quantized_moe_uses_framework_path_and_fails_closed( def framework_benchmark(**kwargs): framework_calls.append(kwargs) - return 1.25, {"power": 100.0}, kernel_source + for num_tokens in kwargs["num_tokens_list"]: + kwargs["on_result"](num_tokens, 1.25, {"power": 100.0}, kernel_source) + return kernel_source run = _load_functions( "run_moe_torch", @@ -781,5 +786,157 @@ def test_framework_moe_source_comes_from_the_constructed_quant_method(): "ModelOptNvFp4FusedMoEMethod", ): assert method_name in framework_bench - assert 'return results["latency_ms"] / len(logits), results["power_stats"], kernel_source' in framework_bench assert 'model_name == "deepseek-ai/DeepSeek-V4-Pro"' not in run_moe + + +def test_moe_population_emits_one_case_per_token_sweep(): + """2026-10-28: SGLang nvfp4 shards took 4.4-5.8 h against 1-1.4 h for vLLM on the same + shapes because every token point was its own case and rebuilt the framework layer. + Cases now carry the whole sweep (vLLM/TRT-LLM layout) and dedup on the series.""" + case = _gptoss_case(tp=1, ep=1) + case.num_tokens_list = [1, 8, 64, 512] + duplicate = SimpleNamespace(**vars(case)) + duplicate.model_name = "example/equivalent-gptoss" + + populated = _populate_gptoss_cases([case, duplicate]) + + assert len(populated) == 1 + assert populated[0][1] == [1, 8, 64, 512] + assert populated[0][8] == "openai/gpt-oss-120b" + + +def test_run_moe_torch_times_every_token_point_on_one_framework_layer(): + logged = [] + framework_calls = [] + fake_torch = SimpleNamespace( + bfloat16=object(), + device=lambda value: value, + set_default_device=lambda _device: None, + cuda=SimpleNamespace( + set_device=lambda _device: None, + empty_cache=lambda: None, + memory_allocated=lambda _device: 0, + get_device_properties=lambda _device: SimpleNamespace(total_memory=100), + get_device_name=lambda _device: "B200", + ), + ) + + def framework_benchmark(**kwargs): + framework_calls.append(kwargs["num_tokens_list"]) + for num_tokens in kwargs["num_tokens_list"]: + kwargs["on_result"](num_tokens, float(num_tokens), {"power": 1.0}, "sglang_flashinfer_trtllm_moe") + return "sglang_flashinfer_trtllm_moe" + + run = _load_functions( + "run_moe_torch", + namespace={ + "torch": fake_torch, + "_benchmark_framework_quantized_moe": framework_benchmark, + "benchmark": lambda *_args, **_kwargs: pytest.fail("nvfp4 must use the framework layer"), + "get_moe_quantization_module_config": lambda *_args, **_kwargs: {}, + "get_sm_version": lambda: 100, + "_fmoe_kernels_mod": SimpleNamespace(_B_DESC_CACHE=SimpleNamespace(clear=lambda: None)), + "gc": SimpleNamespace(collect=lambda: None), + "log_perf": lambda **kwargs: logged.append(kwargs) or True, + "EXIT_CODE_RESTART": 10, + }, + )["run_moe_torch"] + + assert ( + run( + "nvfp4", + [1, 64, 4096], + 7168, + 2048, + 8, + 256, + 1, + 8, + "nvidia/GLM-5-NVFP4", + distributed="power_law", + power_law_alpha=1.2, + moe_backend="flashinfer_trtllm", + perf_filename="moe.csv", + ) + is None + ) + # One layer build for the whole sweep, one persisted row per token point. + assert framework_calls == [[1, 64, 4096]] + assert [row["item_list"][0]["num_tokens"] for row in logged] == [1, 64, 4096] + assert [row["item_list"][0]["latency"] for row in logged] == [1.0, 64.0, 4096.0] + assert {row["item_list"][0]["distribution"] for row in logged} == {"power_law_1.2"} + assert {row["kernel_source"] for row in logged} == {"sglang_flashinfer_trtllm_moe"} + + # A bare int (fills, direct callers) still measures exactly that point. + logged.clear() + run( + "nvfp4", + 768, + 7168, + 2048, + 8, + 256, + 1, + 8, + "nvidia/GLM-5-NVFP4", + moe_backend="flashinfer_trtllm", + perf_filename="moe.csv", + ) + assert [row["item_list"][0]["num_tokens"] for row in logged] == [768] + + +def test_raw_moe_case_persists_each_token_point_and_cleans_up_once(): + logged = [] + bench_calls = [] + cleared = [] + fake_torch = SimpleNamespace( + bfloat16=object(), + device=lambda value: value, + set_default_device=lambda _device: None, + cuda=SimpleNamespace( + set_device=lambda _device: None, + empty_cache=lambda: None, + memory_allocated=lambda _device: 0, + get_device_properties=lambda _device: SimpleNamespace(total_memory=100), + get_device_name=lambda _device: "H200", + ), + ) + + def fake_benchmark(num_tokens, *_args, **_kwargs): + bench_calls.append(num_tokens) + return 0.5 * num_tokens, {"power": 1.0} + + run = _load_functions( + "run_moe_torch", + namespace={ + "torch": fake_torch, + "benchmark": fake_benchmark, + "build_rank0_workloads": lambda **_kwargs: pytest.fail("EP=1 must not build rank-local workloads"), + "get_moe_quantization_module_config": lambda *_args, **_kwargs: {}, + "_benchmark_framework_quantized_moe": lambda **_kwargs: pytest.fail("BF16 must use the raw benchmark"), + "_fmoe_kernels_mod": SimpleNamespace(_B_DESC_CACHE=SimpleNamespace(clear=lambda: cleared.append(True))), + "gc": SimpleNamespace(collect=lambda: None), + "get_sm_version": lambda: 90, + "log_perf": lambda **kwargs: logged.append(kwargs) or True, + "EXIT_CODE_RESTART": 10, + }, + )["run_moe_torch"] + + run( + "bfloat16", + [16, 32], + 4096, + 14336, + 2, + 8, + 1, + 1, + "mistralai/Mixtral-8x7B-v0.1", + distributed="balanced", + perf_filename="moe.csv", + ) + + assert bench_calls == [16, 32] + assert [row["item_list"][0]["num_tokens"] for row in logged] == [16, 32] + assert [row["kernel_source"] for row in logged] == ["sglang_fused_moe_triton"] * 2 + assert cleared == [True] diff --git a/python/aisimulate/tests/unit/collector/sglang/test_collect_moe_unverified_lane_guard.py b/python/aisimulate/tests/unit/collector/sglang/test_collect_moe_unverified_lane_guard.py index 67c911f1d1..f6da27effe 100644 --- a/python/aisimulate/tests/unit/collector/sglang/test_collect_moe_unverified_lane_guard.py +++ b/python/aisimulate/tests/unit/collector/sglang/test_collect_moe_unverified_lane_guard.py @@ -30,12 +30,17 @@ QWEN38MAX_LANES = ["bfloat16", "fp8_block", "nvfp4"] -def _load_guard(installed_version: str): +def _load_guard(installed_version: str, sm_version: int = 90): tree = ast.parse(SOURCE_PATH.read_text(encoding="utf-8"), filename=str(SOURCE_PATH)) function = next( node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "_raise_if_unverified_moe_lane" ) - loaded = {"_dist_version": lambda _name: installed_version, "_check_compat": _check_compat} + loaded = { + "_check_compat": _check_compat, + # the guard reads the runtime through the module helpers; stub both + "_dist_version": lambda _name: installed_version, + "get_sm_version": lambda: sm_version, + } exec(compile(ast.Module(body=[function], type_ignores=[]), str(SOURCE_PATH), "exec"), loaded) return loaded["_raise_if_unverified_moe_lane"] @@ -69,3 +74,30 @@ def test_guard_never_fires_for_qwen38max_collected_lanes(moe_type, installed_ver guard = _load_guard(installed_version) assert guard(moe_type) == installed_version + + +@pytest.mark.parametrize("installed_version", ["0.5.21", "0.5.21+cu130"]) +@pytest.mark.parametrize("moe_type", ["int4_wo", "w4a16_mxfp4"]) +def test_guard_accepts_0521_hopper_lanes_reverified_2026_10_04(moe_type, installed_version): + """0.5.21 Marlin (int4_wo) and Triton (w4a16_mxfp4) dispatch re-verified on SM89/90 only.""" + guard = _load_guard(installed_version, 90) + assert guard(moe_type) == installed_version + + +@pytest.mark.parametrize("moe_type", ["int4_wo", "w4a16_mxfp4", "w4a8_mxfp4_mxfp8"]) +def test_guard_accepts_0521_blackwell_mxfp4_lanes(moe_type): + """SM100/103 MXFP4 lanes re-verified on 0.5.21 from the B200 identity records (6400345f); + int4_wo re-opened 2026-10-06 from the same records (Kimi-K2.5 -> FLASHINFER_TRTLLM / trtllm_gen_moe).""" + guard = _load_guard("0.5.21", 100) + assert guard(moe_type) == "0.5.21" + + +@pytest.mark.parametrize( + "moe_type, sm_version", + [("w4a8_mxfp4_mxfp8", 90), ("int4_wo", 120), ("w4a16_mxfp4", 120)], +) +def test_guard_keeps_0521_unverified_lanes_closed(moe_type, sm_version): + """Every SM120 weight-only lane and the DSV4 FP4 lane on Hopper were not re-verified.""" + guard = _load_guard("0.5.21", sm_version) + with pytest.raises(RuntimeError, match=rf"{moe_type}.*installed: 0.5.21, SM{sm_version}"): + guard(moe_type) diff --git a/python/aisimulate/tests/unit/collector/sglang/test_dsa_forward_tokens.py b/python/aisimulate/tests/unit/collector/sglang/test_dsa_forward_tokens.py index d9eff3a55d..71c8adc05b 100644 --- a/python/aisimulate/tests/unit/collector/sglang/test_dsa_forward_tokens.py +++ b/python/aisimulate/tests/unit/collector/sglang/test_dsa_forward_tokens.py @@ -34,20 +34,20 @@ def test_real_query_tensor_counts_must_match(tokens): input_ids=tensor(tokens), positions=tensor(tokens), out_cache_loc=tensor(tokens), - num_token_non_padded_cpu=tokens, + global_num_token_non_padded_cpu=tokens, ) validate(batch, tokens) batch.input_ids = None - batch.num_token_non_padded_cpu = 0 + batch.global_num_token_non_padded_cpu = 0 with pytest.raises(ValueError, match="input_ids has None"): validate(batch, tokens) -@pytest.mark.parametrize("field", ["positions", "out_cache_loc", "num_token_non_padded_cpu"]) +@pytest.mark.parametrize("field", ["positions", "out_cache_loc", "global_num_token_non_padded_cpu"]) def test_decode_requires_one_new_token_per_request(field): batch = SimpleNamespace( - input_ids=tensor(4), positions=tensor(4), out_cache_loc=tensor(4), num_token_non_padded_cpu=4 + input_ids=tensor(4), positions=tensor(4), out_cache_loc=tensor(4), global_num_token_non_padded_cpu=4 ) - setattr(batch, field, 0 if field == "num_token_non_padded_cpu" else tensor(8192)) + setattr(batch, field, 0 if field == "global_num_token_non_padded_cpu" else tensor(8192)) with pytest.raises(ValueError, match=field): validate(batch, 4) diff --git a/python/aisimulate/tests/unit/collector/sglang/test_dsa_history.py b/python/aisimulate/tests/unit/collector/sglang/test_dsa_history.py index 16ad12aa11..659978b59e 100644 --- a/python/aisimulate/tests/unit/collector/sglang/test_dsa_history.py +++ b/python/aisimulate/tests/unit/collector/sglang/test_dsa_history.py @@ -28,7 +28,7 @@ def test_native_writers_cover_each_history_slot_and_no_new_token(monkeypatch, le torch = pytest.importorskip("torch") if not isinstance(torch, ModuleType) or not isinstance(getattr(torch, "__version__", None), str): pytest.skip("requires real torch, not another collector test's import stub") - quant_module = ModuleType("sglang.srt.layers.attention.dsa.triton_kernel") + quant_module = ModuleType("sglang.kernels.ops.attention.dsa.triton_kernel") quant_module.act_quant = object() monkeypatch.setitem(sys.modules, quant_module.__name__, quant_module) active_context = [] diff --git a/python/aisimulate/tests/unit/collector/sglang/test_dsa_measurement_boundary.py b/python/aisimulate/tests/unit/collector/sglang/test_dsa_measurement_boundary.py index acf9017346..7995cafca6 100644 --- a/python/aisimulate/tests/unit/collector/sglang/test_dsa_measurement_boundary.py +++ b/python/aisimulate/tests/unit/collector/sglang/test_dsa_measurement_boundary.py @@ -17,6 +17,7 @@ from unittest.mock import Mock import pytest +from collector.helper import PhaseClock pytestmark = pytest.mark.unit SOURCE = Path(__file__).resolve().parents[4] / "collector/sglang/collect_mla_module.py" @@ -51,7 +52,7 @@ def set_extend_input_len(self, length): batch = SimpleNamespace(prepare_for_extend=record("prepare")) forward_batch = SimpleNamespace() module( - "sglang.srt.layers.communicator", + "sglang.srt.layers.layer_boundary.adapters.attention", AttentionInputs=lambda *_: object(), get_attn_tp_context=lambda: SimpleNamespace(set_attn_inputs=record("inputs")), ) @@ -128,7 +129,13 @@ def __exit__(self, *_): raise self.close_error graph = GraphContext() - bucket = _load(["graph_token_bucket"], source=GRAPH_SOURCE)["graph_token_bucket"] + + # A runner whose prefill graph captures (4, 144, 2048): the collector keeps SGLang + # 0.5.21 prefill graphs off (test_dsa_prefill_graph.py pins graph_token_bucket), so + # this drives _run_prefill's native-graph branch on its own. + def bucket(_runner, tokens): + return next((size for size in (4, 144, 2048) if size >= tokens), None) + module( "collector.sglang.dsa_prefill_graph", graph_token_bucket=bucket, @@ -155,8 +162,10 @@ def log_perf(**row): "_dsa_skip_indexer_enabled": lambda *_: False, "_resolve_perf_path": lambda _path, filename: filename, "_module_model_architecture": lambda _: "GlmMoeDsaForCausalLM", - "get_version": lambda _: "0.5.14", + "get_version": lambda _: "0.5.21", "log_perf": log_perf, + "PhaseClock": PhaseClock, + "_POINT_PHASE_TOTALS": {}, }, ) kwargs = dict( @@ -174,7 +183,7 @@ def log_perf(**row): model_path="glm", architecture="GlmMoeDsaForCausalLM", backend_name="dsa", - version="0.5.14", + version="0.5.21", device_name="FakeGPU", log_mla_dtype="bfloat16", log_kv_dtype="fp8", diff --git a/python/aisimulate/tests/unit/collector/sglang/test_dsa_prefill_graph.py b/python/aisimulate/tests/unit/collector/sglang/test_dsa_prefill_graph.py index 19867ccf86..a5f8c1a823 100644 --- a/python/aisimulate/tests/unit/collector/sglang/test_dsa_prefill_graph.py +++ b/python/aisimulate/tests/unit/collector/sglang/test_dsa_prefill_graph.py @@ -1,11 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""SGLang prefill uses the native capture buckets, including padded shapes.""" +"""SGLang prefill graph adapter: off at 0.5.21 (eager DSA prefill), closed for any enabled backend.""" import ast +import sys from pathlib import Path -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace import pytest @@ -13,7 +14,7 @@ SOURCE = Path(__file__).resolve().parents[4] / "collector/sglang/dsa_prefill_graph.py" -def select_bucket(backend, sizes, tokens): +def select_bucket(monkeypatch, backend, tokens): node = next( n for n in ast.parse(SOURCE.read_text()).body @@ -21,26 +22,25 @@ def select_bucket(backend, sizes, tokens): ) ns = {} exec(compile(ast.Module(body=[node], type_ignores=[]), str(SOURCE), "exec"), ns) - runner = SimpleNamespace( - server_args=SimpleNamespace( - cuda_graph_config=SimpleNamespace(prefill=SimpleNamespace(backend=backend, bs=sizes)) - ) - ) - return ns["graph_token_bucket"](runner, tokens) + prefill = SimpleNamespace(backend=backend, bs=[1, 128, 256, 2048]) + exec_bag = SimpleNamespace(graph=SimpleNamespace(cuda_graph_config=SimpleNamespace(prefill=prefill))) + for name in ("sglang", "sglang.srt"): + monkeypatch.setitem(sys.modules, name, ModuleType(name)) + runtime_context = ModuleType("sglang.srt.runtime_context") + runtime_context.get_exec = lambda: exec_bag # 0.5.21: the resolved config, not server_args + monkeypatch.setitem(sys.modules, runtime_context.__name__, runtime_context) + return ns["graph_token_bucket"](SimpleNamespace(server_args=SimpleNamespace(cuda_graph_config=None)), tokens) -@pytest.mark.parametrize( - ("tokens", "expected"), [(1, 1), (128, 128), (129, 256), (1057, 2048), (2048, 2048), (2049, None)] -) -def test_prefill_capture_bucket_boundaries(tokens, expected): - assert select_bucket("tc_piecewise", [2048, 128, 1, 256], tokens) == expected +@pytest.mark.parametrize("tokens", [1, 129, 2048, 4096]) +def test_disabled_prefill_graph_times_eagerly(monkeypatch, tokens): + assert select_bucket(monkeypatch, "disabled", tokens) is None -def test_disabled_and_unsupported_backends_are_distinct(): - assert select_bucket("disabled", [1, 128], 1) is None - assert select_bucket("tc_piecewise", [], 1) is None - with pytest.raises(ValueError, match="Unsupported"): - select_bucket("breakable", [1, 128], 1) +@pytest.mark.parametrize("backend", ["tc_piecewise", "breakable", "full"]) +def test_enabled_prefill_graph_fails_closed_until_the_adapter_is_ported(monkeypatch, backend): + with pytest.raises(ValueError, match="not ported to 0.5.21"): + select_bucket(monkeypatch, backend, 1) def test_cleanup_removes_only_new_owned_native_compile_hooks(): diff --git a/python/aisimulate/tests/unit/collector/sglang/test_moe_config_observation.py b/python/aisimulate/tests/unit/collector/sglang/test_moe_config_observation.py index 46edb71f06..0cfbcad28b 100644 --- a/python/aisimulate/tests/unit/collector/sglang/test_moe_config_observation.py +++ b/python/aisimulate/tests/unit/collector/sglang/test_moe_config_observation.py @@ -3,6 +3,7 @@ """The SGLang MoE collector records, never chooses, the Triton config source.""" import json +import re from pathlib import Path from types import SimpleNamespace @@ -85,4 +86,4 @@ def test_sglang_collector_wraps_the_timed_call_and_writes_triton_rows_only(): source = (Path(__file__).resolve().parents[4] / "collector/sglang/collect_moe.py").read_text() assert "with observe_triton_config_selection(_triton_config_namespaces())" in source assert "return [sys.modules[get_config_dtype_str.__module__], globals()]" in source - assert 'if moe_backend == "triton":\n append_triton_config_record(' in source + assert re.search(r'if moe_backend == "triton":\n\s+append_triton_config_record\(', source) diff --git a/python/aisimulate/tests/unit/collector/test_b200_025_publication.py b/python/aisimulate/tests/unit/collector/test_b200_025_publication.py index f177bfcbec..3a320a4ebc 100644 --- a/python/aisimulate/tests/unit/collector/test_b200_025_publication.py +++ b/python/aisimulate/tests/unit/collector/test_b200_025_publication.py @@ -38,7 +38,9 @@ def test_published_025_metadata_loads_and_is_discoverable(): assert event["rows"] == source["source_rows"] assert "classified_failures" not in event assert isinstance(source["classified_failures"], int) and source["classified_failures"] >= 0 - assert total_rows == report["total_published_rows"] == 419538 + # 419538 at publication; the retired dsv4_hca_attn_module table (3918 rows) was + # removed 2026-10-15 and is listed under "withheld" with the reason. + assert total_rows == report["total_published_rows"] == 415620 @pytest.mark.parametrize("field", ["unexpected_runtime_field", "classified_failures"]) diff --git a/python/aisimulate/tests/unit/collector/test_benchmark_graph_scope.py b/python/aisimulate/tests/unit/collector/test_benchmark_graph_scope.py index 8f994eb0e1..7811043d58 100644 --- a/python/aisimulate/tests/unit/collector/test_benchmark_graph_scope.py +++ b/python/aisimulate/tests/unit/collector/test_benchmark_graph_scope.py @@ -6,6 +6,7 @@ import ast import logging import sys +import time from contextlib import contextmanager from functools import partial from pathlib import Path @@ -18,7 +19,7 @@ SOURCE = COLLECTOR / "helper.py" -def trace_benchmark(monkeypatch, *, source=SOURCE, failure=None, explicit=None, eager=False, fallback=False): +def trace_benchmark(monkeypatch, *, source=SOURCE, failure=None, explicit=None, eager=False, fallback=False, **extra): """Execute a helper source snapshot; retain graph references like a traceback.""" events, graphs = [], [] state = {"stream": "caller", "replays": 0, "fault_fired": False} @@ -99,7 +100,13 @@ def elapsed_time(self, end): monkeypatch.setitem(sys.modules, "torch", SimpleNamespace(cuda=cuda)) tree = ast.parse(source.read_text()) node = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "benchmark_with_power") - namespace = {"contextmanager": contextmanager, "get_device_module": lambda device=None: cuda, "logging": logging} + namespace = { + "contextmanager": contextmanager, + "get_device_module": lambda device=None: cuda, + "logging": logging, + "time": time, + "worker_heartbeat": lambda: events.append("heartbeat"), + } exec(compile(ast.Module(body=[node], type_ignores=[]), str(source), "exec"), namespace) kwargs = { "device": SimpleNamespace(index=0), @@ -113,6 +120,7 @@ def elapsed_time(self, end): } if explicit is not None: kwargs["explicit_graph_cleanup"] = explicit + kwargs.update(extra) result, caught = None, None try: with namespace["benchmark_with_power"](**kwargs) as result: @@ -122,7 +130,8 @@ def elapsed_time(self, end): caught = {"type": type(exc).__name__, "message": str(exc), "original_error": exc is error} return { "events": events, - "result": result, + # phase_seconds are wall-clock readings that differ between otherwise identical runs. + "result": {k: v for k, v in result.items() if k != "phase_seconds"} if result is not None else None, "error": caught, "stream": state["stream"], "resets": [graph.resets for graph in graphs], @@ -179,12 +188,32 @@ def test_default_is_identical_to_explicit_false(monkeypatch): assert trace_benchmark(monkeypatch) == trace_benchmark(monkeypatch, explicit=False) +def test_eager_warmups_can_be_fewer_than_replay_warmups(monkeypatch): + observed = trace_benchmark(monkeypatch, num_eager_warmups=1) + assert observed["events"].count("kernel") == 2 # one eager warmup + capture + assert observed["events"].count("graph:replay") == 5 # replay warmups stay at num_warmups (2) + three runs + assert observed["result"]["used_cuda_graph"] is True + zero = trace_benchmark(monkeypatch, num_eager_warmups=0) + assert zero["events"].count("kernel") == 1 # capture only + default = trace_benchmark(monkeypatch) + assert default["events"].count("kernel") == 3 # num_warmups eager passes + capture, unchanged + + +def test_every_measurement_beats_the_stall_watchdog_before_touching_the_kernel(monkeypatch): + graph = trace_benchmark(monkeypatch) + assert graph["events"].count("heartbeat") == 1 + assert graph["events"].index("heartbeat") < graph["events"].index("kernel") + eager = trace_benchmark(monkeypatch, eager=True) + assert eager["events"].count("heartbeat") == 1 + + def test_eager_does_not_touch_graph_lifetime_in_either_mode(monkeypatch): default = trace_benchmark(monkeypatch, eager=True) opted = trace_benchmark(monkeypatch, eager=True, explicit=True) assert default == opted assert default["resets"] == [] - assert default["events"].count("kernel") == 7 + # eager: num_warmups (2) + num_runs (3); the helper no longer repeats the eager warmup after the capture step + assert default["events"].count("kernel") == 5 def test_only_vllm_mla_module_binds_explicit_graph_cleanup(): diff --git a/python/aisimulate/tests/unit/collector/test_collect_ops_not_in_plan_exit_code.py b/python/aisimulate/tests/unit/collector/test_collect_ops_not_in_plan_exit_code.py new file mode 100644 index 0000000000..d0c3ef88f1 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_collect_ops_not_in_plan_exit_code.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``--ops`` naming an op the collector v2 case plan does not contain exits with the +dedicated EXIT_CODE_OPS_NOT_IN_PLAN, not argparse's 2, so a shard scheduler can mark +the shard terminal (sm100 shard plans 2026-10-05: the registry-only vllm +dsv4_hca_attn_module retried to max_attempts on eight systems). That op was removed +on 2026-10-15, so the probe below uses an op another backend's registry owns: a valid +``--ops`` choice that vLLM's case plan never carries, the shape of a stale or +cross-framework shard plan.""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.unit + +PKG_ROOT = Path(__file__).resolve().parents[3] + + +def _run_collect(*args: str) -> subprocess.CompletedProcess: + env = dict(os.environ) + env.pop("COLLECTOR_MODEL_PATH", None) + return subprocess.run( + [sys.executable, str(PKG_ROOT / "collector" / "collect.py"), *args], + cwd=PKG_ROOT, + env=env, + capture_output=True, + text=True, + timeout=300, + ) + + +def test_op_outside_the_backend_plan_exits_with_the_dedicated_code(): + from collector.helper import EXIT_CODE_OPS_NOT_IN_PLAN + + # mamba2 is a trtllm registry op: accepted by --ops, never in vLLM's case plan. + result = _run_collect("--backend", "vllm", "--model-cases-full", "--plan-only", "--ops", "mamba2") + + assert result.returncode == EXIT_CODE_OPS_NOT_IN_PLAN, result.stderr[-2000:] + assert EXIT_CODE_OPS_NOT_IN_PLAN not in (0, 1, 2) + assert "Requested ops are not present in the collector v2 case plan: mamba2" in result.stderr diff --git a/python/aisimulate/tests/unit/collector/test_collect_provenance_writer.py b/python/aisimulate/tests/unit/collector/test_collect_provenance_writer.py index 38affbdca4..efebea3122 100644 --- a/python/aisimulate/tests/unit/collector/test_collect_provenance_writer.py +++ b/python/aisimulate/tests/unit/collector/test_collect_provenance_writer.py @@ -7237,7 +7237,7 @@ def test_dsa_registry_producers_share_transaction_and_checkpoint_ownership(tmp_p tracker = collect_mod._resume_tracker_for_collection( collection, ctx, backend=BACKEND, checkpoint_dir=str(checkpoint_dir), sm_version=100 ) - assert collect_mod._registered_checkpoint_table(tracker._metadata, backend=BACKEND) == table + assert collect_mod._registered_checkpoint_tables(tracker._metadata, backend=BACKEND) == {table} case_id = f"{collection['type']}-case" tracker.mark_attempted(case_id) tracker.mark_passed(case_id) @@ -7420,7 +7420,7 @@ def test_perf_journal_requires_attempts_for_each_table_in_a_batch(monkeypatch, s journal_path = output_root / collect_mod._PERF_TRANSACTION_FILENAME journal = json.loads(journal_path.read_text(encoding="utf-8")) for record in journal["checkpoints"]: - if record["table"] == "mla_bmm_perf": + if "mla_bmm_perf" in record["tables"]: record["attempted"] = [] assert any(record["attempted"] for record in journal["checkpoints"]) with pytest.raises(RuntimeError, match="tables lack attempted checkpoint case IDs"): diff --git a/python/aisimulate/tests/unit/collector/test_collect_replay.py b/python/aisimulate/tests/unit/collector/test_collect_replay.py new file mode 100644 index 0000000000..6c66b576ce --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_collect_replay.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Failure replay (`collect.py --cases-from`): exact-string selection of an earlier run's +failing cases plus a seeded control sample, and the marker that keeps such a run out of +shipped data. Selection semantics only — the executor path is unchanged.""" + +import gzip +import json +import random +import sys +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.unit + +_COLLECTOR_DIR = str(Path(__file__).resolve().parents[3] / "collector") +if _COLLECTOR_DIR not in sys.path: + sys.path.insert(0, _COLLECTOR_DIR) + +import collect + +CASES = [ + [32, 255, 128, 1, 128, False, "flashinfer"], + [32, 511, 128, 1, 128, False, "flashinfer"], + [1, 8191, 96, 1, 128, True, "flashinfer"], + [2, 1023, 32, 8, 128, False, "fa3"], + {"id": "sglang.moe:run:[7, 8]", "params": [7, 8]}, # dict-form task +] + + +def _spec(wanted, control=0): + return collect.ReplaySpec("errors.json", set(wanted), control=control) + + +def test_select_matches_exact_case_string_not_substring(): + spec = _spec({"[32, 511, 128, 1, 128, False, 'flashinfer']"}) + out = spec.select("attention_generation", CASES, random.Random(42)) + assert out == [CASES[1]] + # ", 128, 1, 128, " would substring-match two cases; exact identity keeps one + assert spec.matched["attention_generation"] == {"[32, 511, 128, 1, 128, False, 'flashinfer']"} + assert spec.unmatched() == [] + + +def test_select_uses_params_for_dict_form_tasks(): + spec = _spec({"[7, 8]"}) + out = spec.select("moe", CASES, random.Random(42)) + assert out == [CASES[4]] + + +def test_control_sample_is_seeded_and_excludes_requested(): + wanted = {"[32, 255, 128, 1, 128, False, 'flashinfer']"} + a = _spec(wanted, control=2).select("op", CASES, random.Random(42)) + b = _spec(wanted, control=2).select("op", CASES, random.Random(42)) + assert a == b + assert a[0] == CASES[0] + assert len(a) == 3 and CASES[0] not in a[1:] + big = _spec(wanted, control=100) + assert len(big.select("op", CASES, random.Random(1))) == len(CASES) # capped at the rest + + +def test_load_replay_cases_from_errors_json_summary_gz_and_plain_list(tmp_path: Path): + errors = [ + {"task_params": "[32, 255, 128, 1, 128, False, 'flashinfer']", "error_type": "RuntimeError"}, + {"task_params": None, "error_type": "UnresolvedFailures"}, # checkpoint summary entry, no case + ] + ej = tmp_path / "errors_sglang.attention_generation.json" + ej.write_text(json.dumps(errors)) + assert collect.load_replay_cases(ej) == {"[32, 255, 128, 1, 128, False, 'flashinfer']"} + + summary = tmp_path / "collection_summary_sglang.json.gz" + with gzip.open(summary, "wt") as f: + json.dump({"summary": {"total_errors": 1}, "errors": errors}, f) + assert collect.load_replay_cases(summary) == {"[32, 255, 128, 1, 128, False, 'flashinfer']"} + + plain = tmp_path / "cases.json" + plain.write_text(json.dumps(["[1, 2]", "[3, 4]"])) + assert collect.load_replay_cases(plain) == {"[1, 2]", "[3, 4]"} + + lines = tmp_path / "cases.txt" + lines.write_text("[1, 2]\n\n[3, 4]\n") + assert collect.load_replay_cases(lines) == {"[1, 2]", "[3, 4]"} + + (tmp_path / "empty.json").write_text("[]") + with pytest.raises(ValueError): + collect.load_replay_cases(tmp_path / "empty.json") + + +def test_finish_replay_writes_a_marker_listing_unmatched_requests_across_ops(tmp_path: Path): + spec = _spec({"[32, 255, 128, 1, 128, False, 'flashinfer']", "[9, 9, 9]"}) + spec.select("attention_generation", CASES, random.Random(42)) + spec.select("gemm", [], random.Random(42)) + assert spec.report()["matched_by_op"] == {"attention_generation": 1, "gemm": 0} + import logging + + log = logging.getLogger("replay-test") + collect._finish_replay(spec, tmp_path, log) + marker = json.loads((tmp_path / collect.ReplaySpec.MARKER).read_text()) + assert marker["requested"] == 2 and marker["matched"] == 1 and marker["unmatched"] == ["[9, 9, 9]"] + assert collect._finish_replay(None, tmp_path, log) is None + + +def test_replay_refuses_an_output_root_that_holds_results(tmp_path: Path): + spec = _spec({"[1, 2]"}) + collect._require_replay_output_root_without_results(spec, tmp_path, {}) # empty root: fine + collect._require_replay_output_root_without_results(None, tmp_path, {tmp_path / "gemm_perf.txt": 1}) + with pytest.raises(SystemExit): + collect._require_replay_output_root_without_results(spec, tmp_path, {tmp_path / "gemm_perf.txt": 1}) + (tmp_path / "gemm_perf.parquet").write_bytes(b"") + with pytest.raises(SystemExit): + collect._require_replay_output_root_without_results(spec, tmp_path, {}) diff --git a/python/aisimulate/tests/unit/collector/test_collect_trtllm_alltoall.py b/python/aisimulate/tests/unit/collector/test_collect_trtllm_alltoall.py index 5d25bb7bec..65b565065c 100644 --- a/python/aisimulate/tests/unit/collector/test_collect_trtllm_alltoall.py +++ b/python/aisimulate/tests/unit/collector/test_collect_trtllm_alltoall.py @@ -534,3 +534,81 @@ def test_no_fabricated_node_num_derivation_remains(): # (the docstrings cite the legacy expression, so match its code shape). assert "node_num = max(" not in SOURCE_TEXT assert "world_size // gpus_per_node" in SOURCE_TEXT + + +# --------------------------------------------------------------------------- +# CUDA-graph timing — no eager fallback, capture outcome agreed across ranks +# --------------------------------------------------------------------------- + + +def test_benchmark_op_agrees_capture_before_timing(): + import ast + + tree = ast.parse(SOURCE_TEXT) + function = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "_benchmark_op") + callees = [ast.unparse(node.func) for node in ast.walk(function) if isinstance(node, ast.Call)] + assert callees.index("_agree_graph_capturable") < callees.index("benchmark_with_power") + + +class _Flag: + def __init__(self, value): + self.value = value + + def item(self): + return self.value + + +def _fake_runtime(monkeypatch, *, local_capture_fails, peer_fails, world_size): + import sys + from contextlib import contextmanager + from types import SimpleNamespace + + events = [] + + @contextmanager + def graph(_graph): + events.append("capture") + if local_capture_fails: + raise RuntimeError("operation not permitted when stream is capturing") + yield + + def all_reduce(flag, op=None): + events.append("all_reduce") + flag.value = max(flag.value, 1 if peer_fails else 0) + + dist = SimpleNamespace( + is_available=lambda: True, + is_initialized=lambda: world_size > 1, + get_world_size=lambda: world_size, + all_reduce=all_reduce, + ReduceOp=SimpleNamespace(MAX="max"), + ) + torch = SimpleNamespace( + cuda=SimpleNamespace(CUDAGraph=lambda: object(), graph=graph, synchronize=lambda: events.append("sync")), + tensor=lambda data, **_kwargs: _Flag(data[0]), + int32="int32", + distributed=dist, + ) + monkeypatch.setitem(sys.modules, "torch", torch) + monkeypatch.setitem(sys.modules, "torch.distributed", dist) + return events + + +def test_capture_probe_warms_eagerly_then_captures_and_agrees(monkeypatch): + events = _fake_runtime(monkeypatch, local_capture_fails=False, peer_fails=False, world_size=8) + ata._agree_graph_capturable(lambda: events.append("func"), "dispatch", "cuda:0") + # one eager call (lazy init is not capturable), the probe capture, then the agreement + assert events == ["func", "sync", "capture", "func", "sync", "all_reduce"] + + +def test_local_capture_failure_is_agreed_then_raises_classified(monkeypatch): + events = _fake_runtime(monkeypatch, local_capture_fails=True, peer_fails=False, world_size=8) + with pytest.raises(ata.TrtllmAlltoallBenchmarkError, match=r"CUDA graph capture of combine failed \(RuntimeError"): + ata._agree_graph_capturable(lambda: None, "combine", "cuda:0") + assert "all_reduce" in events # the peers learn before any of them replays + + +def test_peer_capture_failure_raises_on_this_rank_too(monkeypatch): + _fake_runtime(monkeypatch, local_capture_fails=False, peer_fails=True, world_size=8) + with pytest.raises(ata.TrtllmAlltoallBenchmarkError, match="peer rank"): + ata._agree_graph_capturable(lambda: None, "prepare", "cuda:0") diff --git a/python/aisimulate/tests/unit/collector/test_dsv4_common_test_cases.py b/python/aisimulate/tests/unit/collector/test_dsv4_common_test_cases.py index b786a15261..fea3f3b580 100644 --- a/python/aisimulate/tests/unit/collector/test_dsv4_common_test_cases.py +++ b/python/aisimulate/tests/unit/collector/test_dsv4_common_test_cases.py @@ -127,7 +127,7 @@ def test_dsv4_cases_skip_unrelated_model_filter(monkeypatch): monkeypatch.setattr(sys, "argv", ["pytest"]) assert common_test_cases.get_dsv4_csa_context_test_cases() == [] - assert common_test_cases.get_dsv4_hca_attn_test_cases() == [] + assert common_test_cases.get_dsv4_paged_mqa_logits_test_cases() == [] def test_sglang_sparse_modules_topk_calib_uses_canonical_model(monkeypatch, capsys): diff --git a/python/aisimulate/tests/unit/collector/test_error_classification.py b/python/aisimulate/tests/unit/collector/test_error_classification.py new file mode 100644 index 0000000000..fd6d11263b --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_error_classification.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A collector-raised FIXME(kernel-limit) failure is reported as ``kernel_limit``, +everything else as ``unexpected`` (pipeline summaries classified by hand until +2026-10-08). Runs with PYTHONPATH=collector like test_registry_table_producers.""" + +import collect +import pytest + +pytestmark = pytest.mark.unit + + +def test_kernel_limit_prefix_is_classified(): + assert collect._classify_exception(RuntimeError("FIXME(kernel-limit): no head_dim=192 kernel")) == "kernel_limit" + assert collect._classify_exception(ValueError(" FIXME(kernel-limit): divisibility")) == "kernel_limit" + assert collect._classify_exception(RuntimeError("CUDA error: an illegal memory access")) == "unexpected" + assert collect._classify_exception(AssertionError()) == "unexpected" diff --git a/python/aisimulate/tests/unit/collector/test_framework_manifest.py b/python/aisimulate/tests/unit/collector/test_framework_manifest.py index d37fff566c..8b6d5a28a1 100644 --- a/python/aisimulate/tests/unit/collector/test_framework_manifest.py +++ b/python/aisimulate/tests/unit/collector/test_framework_manifest.py @@ -15,7 +15,6 @@ ) from collector.sglang.registry import REGISTRY as SGLANG_REGISTRY from collector.trtllm.registry import REGISTRY as TRTLLM_REGISTRY -from collector.version_resolver import _check_compat from collector.vllm.registry import REGISTRY as VLLM_REGISTRY from collector.vllm.registry import REGISTRY_XPU as VLLM_XPU_REGISTRY from collector.wideep.sglang import dataset_version_label @@ -33,14 +32,13 @@ def test_manifest_exposes_current_framework_versions_and_images(): trtllm = get_collector_runtime("trtllm") vllm = get_collector_runtime("vllm") - assert sglang.version == "0.5.14" - assert sglang.image().startswith("lmsysorg/sglang:v0.5.14@sha256:") - assert sglang.image("cu130").startswith("lmsysorg/sglang:v0.5.14-cu130@sha256:") - assert trtllm.version == "1.3.0rc20" - assert trtllm.image().startswith("nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20@sha256:") - assert vllm.version == "0.24.0" - assert vllm.image().startswith("vllm/vllm-openai:v0.24.0@sha256:") - assert vllm.image("cu129").startswith("vllm/vllm-openai:v0.24.0-cu129@sha256:") + assert sglang.version == "0.5.21" + assert sglang.image().startswith("lmsysorg/sglang:v0.5.21@sha256:") + assert sglang.image("cu130").startswith("lmsysorg/sglang:v0.5.21-cu130@sha256:") + assert trtllm.version == "1.3.0rc29" + assert trtllm.image().startswith("nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc29@sha256:") + assert vllm.version == "0.30.0" + assert vllm.image().startswith("vllm/vllm-openai:v0.30.0@sha256:") # Unknown variants intentionally fall back to the pinned default image. assert vllm.image("cu130") == vllm.image() @@ -57,11 +55,9 @@ def test_manifest_exposes_pinned_vllm_xpu_runtime_identity(): def test_active_cuda_vllm_collectors_are_exactly_pinned_to_manifest_version(): assert all(not entry.versions for entry in VLLM_REGISTRY) - # Each module pins the runtime that actually collects it: the manifest - # default, or its family override (e.g. kda runs only on the vllm kimi-k3 - # preview image, frameworks.vllm.families.kda). The Qwen3.8 target-lane - # collectors accept both the default 0.24.0 runtime and model-pinned - # 0.27.1, so compare the declaration semantically just as collect.py does. + # Each module pins exactly the runtime that collects it: the manifest + # default, or its family override (frameworks.vllm.families, e.g. gemm on + # 0.31.0). One release per file (owner policy 2026-09-20): no ranges. module_versions: dict[str, set[str]] = {} for entry in VLLM_REGISTRY: module_versions.setdefault(entry.module, set()).add(resolve_op_runtime("vllm", entry.op).version) @@ -72,22 +68,36 @@ def test_active_cuda_vllm_collectors_are_exactly_pinned_to_manifest_version(): source = (REPO_ROOT / f"{module.replace('.', '/')}.py").read_text(encoding="utf-8") declarations = [line.strip() for line in source.splitlines() if line.startswith("__compat__")] assert len(declarations) == 1, module - declared = declarations[0].split("=", 1)[1].strip().strip('"') - assert _check_compat(declared, resolved_version), (module, declared, resolved_version) + assert declarations[0] == f'__compat__ = "vllm=={resolved_version}"', (module, declarations[0]) -@pytest.mark.parametrize( - "module", - ["collector.vllm.collect_moe", "collector.vllm.collect_gdn", "collector.vllm.collect_gemm"], -) -def test_vllm_target_lane_collectors_declare_the_exact_bumped_compat_range(module): - expected = '__compat__ = "vllm>=0.24.0,<=0.27.1,!=0.25.1,!=0.26.0,!=0.27.0"' - if module in {"collector.vllm.collect_gemm", "collector.vllm.collect_moe"}: - # GLM-5.3-Flash model pin: 0.31.0 is admitted; 0.28.0-0.30.0 are not. - expected = '__compat__ = "vllm>=0.24.0,<=0.31.0,!=0.25.1,!=0.26.0,!=0.27.0,!=0.28.0,!=0.29.0,!=0.30.0"' - source = (REPO_ROOT / f"{module.replace('.', '/')}.py").read_text(encoding="utf-8") - declarations = [line.strip() for line in source.splitlines() if line.startswith("__compat__")] - assert declarations == [expected], module +# Known exception, not a pattern: collect_dsv41_module carries the pre-existing +# source-build pin from #160 (sglang@1aa0e962...). framework_manifest.yaml never +# declares that build, so the module cannot run on any manifest-resolved runtime; +# fixing the pin is out of scope here. Any other `@` pin fails the test below. +_SOURCE_BUILD_PIN_EXCEPTIONS = { + "collector.sglang.collect_dsv41_module": '__compat__ = "sglang@1aa0e962b206102b7c439a4a0c4981cfec6e87bc"', +} + + +@pytest.mark.parametrize("backend", ["sglang", "trtllm"]) +def test_stock_collectors_are_exactly_pinned_to_their_resolved_runtime(backend): + # One release per file (owner policy 2026-09-20): every registry module + # declares exactly the runtime the manifest resolves for its op (default or + # family pin). Earlier releases' code is `git log` away, never a range. + registry = {"sglang": SGLANG_REGISTRY, "trtllm": TRTLLM_REGISTRY}[backend] + module_versions: dict[str, set[str]] = {} + for entry in registry: + module_versions.setdefault(entry.module, set()).add(resolve_op_runtime(backend, entry.op).version) + for module, versions in sorted(module_versions.items()): + assert len(versions) == 1, (module, versions) + source = (REPO_ROOT / f"{module.replace('.', '/')}.py").read_text(encoding="utf-8") + declarations = [line.strip() for line in source.splitlines() if line.startswith("__compat__")] + version = next(iter(versions)) + allowed = {f'__compat__ = "{backend}=={version}"'} + if module in _SOURCE_BUILD_PIN_EXCEPTIONS: + allowed = {_SOURCE_BUILD_PIN_EXCEPTIONS[module]} + assert len(declarations) == 1 and declarations[0] in allowed, (module, declarations) def test_active_vllm_xpu_collectors_are_exactly_pinned_to_manifest_version(): @@ -303,11 +313,11 @@ def test_wideep_entry_missing_data_backend_is_rejected(tmp_path): @pytest.mark.parametrize( ("installed_version", "requested_ops", "workload", "version"), [ - # "all ops" is no longer resolvable in one container for sglang — the - # kda family pins the kimi-k3 branch runtime (0.5.16), so the default - # expectation is asserted on an explicit default-family op instead. - ("0.5.14+cu130", {"gemm"}, "default", "0.5.14"), - ("0.5.16", {"kda"}, "default", "0.5.16"), + # Since the 0.5.21 default (2026-10-01) kda and msa resolve on the + # stock runtime again (the kimi-k3 / v0.5.16 family pins are retired). + ("0.5.21+cu130", {"gemm"}, "default", "0.5.21"), + ("0.5.21", {"kda"}, "default", "0.5.21"), + ("0.5.21", {"msa_context_module"}, "default", "0.5.21"), ("0.5.10", {"moe_ep"}, "wideep", "0.5.10"), ], ) @@ -332,14 +342,11 @@ def test_vllm_xpu_runtime_selection_rejects_version_mismatch(): @pytest.mark.parametrize( ("installed_version", "requested_ops", "match"), [ - ("0.5.13", {"gemm"}, r"stock collector requires exactly 0\.5\.14"), - ("0.5.14rc1", {"gemm"}, r"stock collector requires exactly 0\.5\.14"), - ("0.5.14.post1", {"gemm"}, r"stock collector requires exactly 0\.5\.14"), - ("0.5.14", {"moe_ep"}, r"WideEP collector requires exactly 0\.5\.10"), - ("0.5.14", {"gemm", "moe_ep"}, r"0\.5\.14 != 0\.5\.10.*separate containers"), - # kda runs only on the kimi-k3 branch runtime (families.kda pin): - # mixing it with a default-family op must fail closed. - ("0.5.14", {"gemm", "kda"}, r"multiple runtime versions"), + ("0.5.14", {"gemm"}, r"stock collector requires exactly 0\.5\.21"), + ("0.5.21rc1", {"gemm"}, r"stock collector requires exactly 0\.5\.21"), + ("0.5.21.post1", {"gemm"}, r"stock collector requires exactly 0\.5\.21"), + ("0.5.21", {"moe_ep"}, r"WideEP collector requires exactly 0\.5\.10"), + ("0.5.21", {"gemm", "moe_ep"}, r"0\.5\.21 != 0\.5\.10.*separate containers"), ], ) def test_runtime_selection_rejects_mismatched_or_mixed_pins(installed_version, requested_ops, match): @@ -353,8 +360,8 @@ def test_runtime_selection_rejects_mismatched_or_mixed_pins(installed_version, r @pytest.mark.parametrize( ("installed_version", "requested_ops", "workload", "version"), [ - ("0.5.14+cu130", {"gemm"}, "default", "0.5.14"), - ("0.5.16", {"kda"}, "default", "0.5.16"), + ("0.5.21+cu130", {"gemm"}, "default", "0.5.21"), + ("0.5.21", {"kda"}, "default", "0.5.21"), ("0.5.10", {"moe_ep"}, "wideep", "0.5.10"), ], ) @@ -374,10 +381,10 @@ def test_no_model_identity_resolves_exactly_like_today(installed_version, reques def test_unknown_model_id_falls_back_to_default_resolution(): - baseline = require_collector_runtime("sglang", "0.5.14", requested_ops={"gemm"}, wideep_ops=WIDEEP_OPS) + baseline = require_collector_runtime("sglang", "0.5.21", requested_ops={"gemm"}, wideep_ops=WIDEEP_OPS) unmatched = require_collector_runtime( "sglang", - "0.5.14", + "0.5.21", requested_ops={"gemm"}, wideep_ops=WIDEEP_OPS, model_path="some-org/not-a-pinned-model", @@ -385,81 +392,26 @@ def test_unknown_model_id_falls_back_to_default_resolution(): assert unmatched == baseline -@pytest.mark.parametrize( - "model_path", - [ - "Qwen/Qwen3.8-2.4T-A95B", - "Qwen/Qwen3.8-2.4T-A95B-FP8", - "RadixArk/Qwen3.8-2.4T-A95B-NVFP4", - ], -) -def test_model_pin_match_resolves_qwen38_max_to_sglang_0_5_17(model_path): - # Real manifest (no path= override): the checked-in frameworks.sglang.models - # entries added for Qwen3.8-Max day-0 support. - runtime = require_collector_runtime( - "sglang", "0.5.17", requested_ops={"gemm"}, wideep_ops=WIDEEP_OPS, model_path=model_path - ) - assert runtime.version == "0.5.17" - assert runtime.image().startswith("lmsysorg/sglang:v0.5.17@sha256:") - assert runtime.image("cu130").startswith("lmsysorg/sglang:v0.5.17-cu130@sha256:") - # Not a family classification — see _model_pinned_runtime docstring. - assert runtime.family is None - - -def test_model_pin_mismatch_error_names_the_model_scoped_image(): - with pytest.raises(RuntimeError) as excinfo: - require_collector_runtime( - "sglang", - "0.5.14", - requested_ops={"gemm"}, - wideep_ops=WIDEEP_OPS, - model_path="Qwen/Qwen3.8-2.4T-A95B", - ) - message = str(excinfo.value) - # Same template as the pre-4b guard ("~:249"), but naming the - # model-scoped runtime/image instead of the framework default. - assert "sglang stock collector requires exactly 0.5.17, found 0.5.14" in message - assert "use lmsysorg/sglang:v0.5.17@sha256:" in message - - -@pytest.mark.parametrize( - "model_path", - [ - "Qwen/Qwen3.8-2.4T-A95B", - "Qwen/Qwen3.8-2.4T-A95B-FP8", - "RadixArk/Qwen3.8-2.4T-A95B-NVFP4", - ], -) -def test_model_pin_match_resolves_qwen38_max_to_vllm_0_27_1(model_path): - runtime = require_collector_runtime( - "vllm", "0.27.1", requested_ops={"gemm"}, wideep_ops=set(), model_path=model_path - ) - assert runtime.version == "0.27.1" - assert runtime.image().startswith("vllm/vllm-openai:v0.27.1@sha256:") - assert runtime.image("cu129").startswith("vllm/vllm-openai:v0.27.1-cu129@sha256:") - assert runtime.family is None - - def test_vllm_model_pin_mismatch_error_names_the_model_scoped_image(): with pytest.raises(RuntimeError) as excinfo: require_collector_runtime( "vllm", - "0.24.0", - requested_ops={"gemm"}, + "0.30.0", + requested_ops={"attention_context"}, wideep_ops=set(), - model_path="Qwen/Qwen3.8-2.4T-A95B", + model_path="zai-org/GLM-5.3-Flash", ) message = str(excinfo.value) - assert "vllm stock collector requires exactly 0.27.1, found 0.24.0" in message - assert "use vllm/vllm-openai:v0.27.1@sha256:" in message + assert "vllm stock collector requires exactly 0.31.0, found 0.30.0" in message + assert "use vllm/vllm-openai:v0.31.0@sha256:" in message def test_vllm_unknown_model_id_falls_back_to_default_resolution(): - baseline = require_collector_runtime("vllm", "0.24.0", requested_ops={"gemm"}, wideep_ops=set()) + baseline = require_collector_runtime("vllm", "0.30.0", requested_ops={"attention_context"}, wideep_ops=set()) unmatched = require_collector_runtime( "vllm", - "0.24.0", - requested_ops={"gemm"}, + "0.30.0", + requested_ops={"attention_context"}, wideep_ops=set(), model_path="some-org/not-a-pinned-model", ) @@ -476,7 +428,7 @@ def test_real_manifest_models_section_does_not_break_validate_resolution(): def test_unknown_requested_op_fails_with_key_error(): with pytest.raises(KeyError, match=r"has no op\(s\): \['not_a_real_op'\]"): - require_collector_runtime("sglang", "0.5.14", requested_ops={"not_a_real_op"}, wideep_ops=set()) + require_collector_runtime("sglang", "0.5.21", requested_ops={"not_a_real_op"}, wideep_ops=set()) def test_vllm_xpu_unknown_requested_op_fails_with_key_error(): @@ -487,7 +439,7 @@ def test_vllm_xpu_unknown_requested_op_fails_with_key_error(): def test_typo_mixed_with_real_op_fails_closed(): # A typo must not be silently dropped just because another requested op is valid. with pytest.raises(KeyError, match=r"has no op\(s\): \['not_a_real_op'\]"): - require_collector_runtime("sglang", "0.5.14", requested_ops={"gemm", "not_a_real_op"}, wideep_ops=set()) + require_collector_runtime("sglang", "0.5.21", requested_ops={"gemm", "not_a_real_op"}, wideep_ops=set()) def test_wideep_registry_entries_are_separate_from_stock_backend_registries(): @@ -787,34 +739,6 @@ def test_model_pin_images_must_be_digest_pinned(tmp_path): get_collector_runtime("sglang", path=manifest) -@pytest.mark.parametrize( - "version,accepted", - [ - ("0.24.0", True), - ("0.25.0", True), - ("0.25.1", False), - ("0.26.0", False), - ("0.27.0", False), - ("0.27.1", True), - ("0.28.0", False), - ("0.29.0", False), - ("0.30.0", False), - ("0.31.0", True), - ("0.31.1", False), - ], -) -def test_gemm_025_qualification_preserves_other_release_gaps(version, accepted): - import ast - - source = ast.parse((COLLECTOR_ROOT / "vllm" / "collect_gemm.py").read_text()) - declaration = next( - node.value.value - for node in source.body - if isinstance(node, ast.Assign) and any(isinstance(t, ast.Name) and t.id == "__compat__" for t in node.targets) - ) - assert _check_compat(declaration, version) is accepted - - @pytest.fixture def explicit_vllm_025_manifest(tmp_path): """Synthetic loader fixture, not a shipped image or collection declaration.""" @@ -896,4 +820,4 @@ def test_explicit_nondefault_manifest_path_is_not_overridden(tmp_path, monkeypat path.write_bytes((COLLECTOR_ROOT / "framework_manifest.yaml").read_bytes()) monkeypatch.setenv(RUNTIME_MANIFEST_ENV, "/missing/ambient/manifest") monkeypatch.setenv(RUNTIME_MANIFEST_SHA256_ENV, "0" * 64) - assert get_collector_runtime("vllm", path=path).version == "0.24.0" + assert get_collector_runtime("vllm", path=path).version == "0.30.0" diff --git a/python/aisimulate/tests/unit/collector/test_fullnode.py b/python/aisimulate/tests/unit/collector/test_fullnode.py index 00b5d519cc..f0dea2cc2a 100644 --- a/python/aisimulate/tests/unit/collector/test_fullnode.py +++ b/python/aisimulate/tests/unit/collector/test_fullnode.py @@ -172,6 +172,71 @@ def test_runner_fails_closed_when_failure_list_is_missing(monkeypatch): assert "must return" in errors[0]["error_message"] +def test_replay_selects_full_node_cases_like_the_pool_ops(monkeypatch): + cases = [{"shape": 0}, {"shape": 1}, {"shape": 2}] + ran = [] + + def run_cases(*, perf_filename, limit, cases): + ran.extend(cases) + return {"succeeded": len(cases), "failed": []} + + module_name = "fake_fullnode_collector_replay" + monkeypatch.setitem( + sys.modules, + module_name, + SimpleNamespace(__compat__=">=0.5.0", get_cases=lambda: list(cases), run_cases=run_cases), + ) + monkeypatch.setattr(fullnode, "filter_cases", lambda values, **_kwargs: (values, [])) + replay = collect_mod.ReplaySpec("errors.json", {str(cases[1])}) + + errors = fullnode.collect_sglang_fullnode_op( + { + "name": "fake", + "type": "deepep_ll", + "module": module_name, + "get_func": "get_cases", + "run_func": "run_cases", + "perf_filename": "perf.txt", + }, + runtime_version="0.5.12", + limit=None, + shuffle=False, + shuffle_seed=42, + backend="sglang", + resume_options=None, + model_path=None, + case_plan=None, + sm_version=100, + case_filters=None, + replay=replay, + get_test_cases_for_model=lambda get_func, _model_path: get_func(), + resume_checkpoint_cls=_Checkpoint, + logger=_logger(), + ) + + assert errors == [] + assert ran == [cases[1]] + assert replay.report()["matched_by_op"] == {"deepep_ll": 1} + + +def test_collect_sglang_passes_only_keywords_the_full_node_runner_accepts(): + """collect_sglang runs full-node ops after the pool ops: a keyword the runner does not + accept is a TypeError that loses the run's summary and provenance.""" + import ast + import inspect + + tree = ast.parse(Path(collect_mod.__file__).read_text(encoding="utf-8")) + calls = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) and getattr(node.func, "id", None) == "collect_sglang_fullnode_op" + ] + assert calls + accepted = set(inspect.signature(fullnode.collect_sglang_fullnode_op).parameters) + for call in calls: + assert {keyword.arg for keyword in call.keywords} <= accepted + + REAL_FULLNODE_MODULE = "collector.wideep.sglang.collect_deepep_ll" diff --git a/python/aisimulate/tests/unit/collector/test_getter_deduplication.py b/python/aisimulate/tests/unit/collector/test_getter_deduplication.py index a3c9d87ffc..c37078e5bf 100644 --- a/python/aisimulate/tests/unit/collector/test_getter_deduplication.py +++ b/python/aisimulate/tests/unit/collector/test_getter_deduplication.py @@ -137,6 +137,7 @@ def _install_vllm_stubs(monkeypatch): "collector.helper", balanced_logits=_noop, benchmark_with_power=_noop, + flashinfer_autotuned=_noop, get_sm_version=lambda: 100, log_perf=_noop, power_law_logits_v3=_noop, diff --git a/python/aisimulate/tests/unit/collector/test_glm5next_collection_plan.py b/python/aisimulate/tests/unit/collector/test_glm5next_collection_plan.py index 6c98432a6a..d120d01013 100644 --- a/python/aisimulate/tests/unit/collector/test_glm5next_collection_plan.py +++ b/python/aisimulate/tests/unit/collector/test_glm5next_collection_plan.py @@ -3,10 +3,7 @@ """GLM-5.3-Flash standalone per-op collection: case plan, routing and version gates.""" -import ast import json -import sys -import types from pathlib import Path import pytest @@ -18,10 +15,8 @@ moe_model_allows_quantization, ) from collector.model_cases import build_collection_case_plan -from collector.sglang.registry import REGISTRY as SGLANG_REGISTRY -from collector.version_resolver import _check_compat -from .test_getter_deduplication import _install_vllm_stubs, _load_collector, _stub_module +from .test_getter_deduplication import _install_vllm_stubs, _load_collector pytestmark = pytest.mark.unit @@ -33,22 +28,20 @@ GLM_OFFGRID_WIDTHS = {8, 16, 32, 288, 3336, 6416, 12576, 19360, 24576, 24896, 38720, 77440, 154880} -def _compat(relative_path: str) -> str: - tree = ast.parse((REPO_ROOT / relative_path).read_text()) - return next( - node.value.value - for node in tree.body - if isinstance(node, ast.Assign) and any(getattr(t, "id", None) == "__compat__" for t in node.targets) - ) - - -@pytest.mark.parametrize("backend", ["vllm", "sglang"]) +@pytest.mark.parametrize( + ("backend", "ops"), + [ + ("vllm", {"gemm", "moe", "kda", "mhc_module"}), + # SGLang 0.5.21 is not an audited GLM KDA / mHC runtime: scheduling them only logs failures + ("sglang", {"gemm", "moe"}), + ], +) @pytest.mark.parametrize("model_path", [FP8, NVFP4]) -def test_plan_collects_gemm_moe_kda_and_mhc(backend, model_path): +def test_plan_collects_the_audited_glm_ops(backend, ops, model_path): plan = build_collection_case_plan(backend=backend, model_path=model_path, gpu_type="gb300") assert plan.model_architecture == "Glm5NextForConditionalGeneration" assert plan.sm_version == 103 - assert set(plan.ops) == {"gemm", "moe", "kda", "mhc_module"} + assert set(plan.ops) == ops def test_trtllm_plan_has_no_glm_moe_lane(): @@ -162,89 +155,6 @@ def test_vllm_nvfp4_uses_modelopt_only_for_the_glm53flash_checkpoint(monkeypatch assert (qc["quant_method"], qc["quant_algo"]) == ("modelopt", "NVFP4") -@pytest.mark.parametrize( - "path", - ["collector/vllm/collect_gemm.py", "collector/vllm/collect_moe.py"], -) -@pytest.mark.parametrize( - "version,accepted", - [ - ("0.31.0", True), - ("0.30.0", False), - ("0.28.0", False), - ("0.29.0", False), - ("0.31.1", False), - ("0.32.0", False), - ("0.24.0", True), - ], -) -def test_vllm_collectors_admit_exactly_the_audited_releases(path, version, accepted): - assert _check_compat(_compat(path), version) is accepted - - -@pytest.mark.parametrize( - "path,fork,names", - [("collector/sglang/collect_gemm.py", "collect_gemm_0520", ("get_gemm_test_cases", "run_gemm"))], -) -def test_sglang_0_5_20_runs_delegate_to_the_serving_built_fork(path, fork, names): - """The registry module stays unversioned (changed_ops rejects VersionRoute - entries) and swaps in the 0.5.20 implementation by installed release.""" - tree = ast.parse((REPO_ROOT / path).read_text()) - gate = next( - node - for node in tree.body - if isinstance(node, ast.If) - and "_sglang_release()" in ast.unparse(node.test) - and "(0, 5, 20)" in ast.unparse(node.test) - ) - body = ast.unparse(gate) - assert f"from collector.sglang import {fork} as _fork_0520" in body - for name in names: - assert f"{name} = _fork_0520.{name}" in body - entry = next(e for e in SGLANG_REGISTRY if e.module == f"collector.sglang.{Path(path).stem}") - assert not entry.versions - compat = _compat(path) - assert _check_compat(compat, "0.5.20") and _check_compat(compat, "0.5.14") - for version in ("0.5.18", "0.5.19"): - assert not _check_compat(compat, version) - assert _check_compat(_compat(f"collector/sglang/{fork}.py"), "0.5.20") - assert not _check_compat(_compat(f"collector/sglang/{fork}.py"), "0.5.17") - - -def test_sglang_moe_admits_0_5_20_but_not_unaudited_releases(): - compat = _compat("collector/sglang/collect_moe.py") - for version, accepted in [ - ("0.5.14", True), - ("0.5.17", True), - ("0.5.18", False), - ("0.5.19", False), - ("0.5.20", True), - ]: - assert _check_compat(compat, version) is accepted, version - - -def _load_sglang_gemm_0520(monkeypatch): - _stub_module(monkeypatch, "torch", bfloat16="bfloat16", float32="float32") - helper = types.ModuleType("collector.helper") - helper.benchmark_with_power = None - helper.log_perf = None - helper.get_sm_version = lambda: 103 - monkeypatch.setitem(sys.modules, "collector.helper", helper) - return _load_collector(monkeypatch, "collector.sglang.collect_gemm_0520", "collector/sglang/collect_gemm_0520.py") - - -def test_sglang_0520_queues_small_shapes_the_v1_fixme_skipped(monkeypatch): - monkeypatch.setenv("COLLECTOR_MODEL_PATH", FP8) - module = _load_sglang_gemm_0520(monkeypatch) - cases = module.get_gemm_test_cases() - types_by_shape = {} - for gemm_type, x, n, k in cases: - types_by_shape.setdefault((n, k), set()).add(gemm_type) - # Per-tensor fp8 has no 0.5.20 consumer and is not planned. - assert types_by_shape[(288, 4096)] == {"bfloat16", "fp8_block", "nvfp4"} - assert types_by_shape[(32, 4096)] == {"bfloat16", "fp8_block", "nvfp4"} - - def test_sglang_nvfp4_allowlist_admits_glm53flash_nvfp4_only(): import yaml diff --git a/python/aisimulate/tests/unit/collector/test_glm5next_runtime_pins.py b/python/aisimulate/tests/unit/collector/test_glm5next_runtime_pins.py index 1e5c52fee3..fabaa97cfa 100644 --- a/python/aisimulate/tests/unit/collector/test_glm5next_runtime_pins.py +++ b/python/aisimulate/tests/unit/collector/test_glm5next_runtime_pins.py @@ -1,7 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GLM-5.3-Flash model-scoped collector runtime pins (stock vLLM 0.31.0, SGLang 0.5.20).""" +"""GLM-5.3-Flash model-scoped collector runtime pins (stock vLLM 0.31.0; the SGLang 0.5.20 pin +retired 2026-10-30, one release per collector file).""" import pytest from collector.framework_manifest import require_collector_runtime, validate_resolution @@ -10,7 +11,6 @@ GLM_MODEL_PATHS = ("zai-org/GLM-5.3-Flash", "nvidia/GLM-5.3-Flash-NVFP4") VLLM_INDEX = "vllm/vllm-openai:v0.31.0@sha256:c1c9f6fd5c109ba7f0546a59f5b2f15fb87f64c77782e90a27b648b42a8e67c3" -SGLANG_INDEX = "lmsysorg/sglang:v0.5.20@sha256:06e4f2ed21afde4ff513cda65070124e727ba23ccaeff7712b8c40e1097d611f" @pytest.mark.parametrize("model_path", GLM_MODEL_PATHS) @@ -26,20 +26,9 @@ def test_glm_vllm_pin_is_the_stock_0_31_0_release(model_path): assert runtime.family is None -@pytest.mark.parametrize("model_path", GLM_MODEL_PATHS) -def test_glm_sglang_pin_resolves_every_op_to_0_5_20(model_path): - runtime = require_collector_runtime( - "sglang", "0.5.20", requested_ops={"gemm", "moe", "kda", "mhc_module"}, model_path=model_path - ) - assert runtime.version == "0.5.20" - assert runtime.image() == SGLANG_INDEX - assert runtime.image("cu130").endswith(SGLANG_INDEX.split("@", 1)[1]) - assert runtime.source_commit == "94602c9c2b7cbdb8efd5c52802dac6a1c180089e" - - @pytest.mark.parametrize( "backend,wrong_version,expected", - [("vllm", "0.27.1", "0.31.0"), ("vllm", "0.30.0", "0.31.0"), ("sglang", "0.5.17", "0.5.20")], + [("vllm", "0.27.1", "0.31.0"), ("vllm", "0.30.0", "0.31.0")], ) def test_glm_pin_rejects_other_runtimes(backend, wrong_version, expected): with pytest.raises(RuntimeError, match=f"requires exactly {expected}"): @@ -49,12 +38,9 @@ def test_glm_pin_rejects_other_runtimes(backend, wrong_version, expected): @pytest.mark.parametrize( "backend,version,model_path", [ - ("vllm", "0.24.0", None), - ("vllm", "0.24.0", "zai-org/GLM-5.2-FP8"), - ("vllm", "0.27.1", "Qwen/Qwen3.8-2.4T-A95B"), - ("sglang", "0.5.14", None), - ("sglang", "0.5.14", "zai-org/GLM-5.2-FP8"), - ("sglang", "0.5.17", "Qwen/Qwen3.8-2.4T-A95B"), + ("vllm", "0.31.0", None), # gemm family pin + ("vllm", "0.31.0", "zai-org/GLM-5.2-FP8"), + ("sglang", "0.5.21", "zai-org/GLM-5.2-FP8"), ], ) def test_other_models_keep_their_runtime(backend, version, model_path): diff --git a/python/aisimulate/tests/unit/collector/test_helper_flashinfer_autotune.py b/python/aisimulate/tests/unit/collector/test_helper_flashinfer_autotune.py new file mode 100644 index 0000000000..f427d9deb6 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_helper_flashinfer_autotune.py @@ -0,0 +1,226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``helper.flashinfer_autotuned`` times FlashInfer-backed kernels on the autotuned +tactic (serving truth) instead of FlashInfer's untuned fallback, whose default kernel +moved between FlashInfer 0.6.12 and 0.6.18 and broke the SM100 nvfp4 MoE EP sweeps of +the vLLM 0.30.0 / SGLang 0.5.21 campaign tables (2026-10-27).""" + +from __future__ import annotations + +import sys +import types +from contextlib import contextmanager + +import pytest +from collector import helper +from collector.helper import ( + FLASHINFER_AUTOTUNE_BUDGET_ENV, + FLASHINFER_AUTOTUNE_ENV, + flashinfer_autotuned, +) + +pytestmark = pytest.mark.unit + + +class _FakeTuner: + instance = None + + def __init__(self): + self.profiling_cache = {} + # FlashInfer's own per-tactic profiling budget (AutoTuner.__init__ defaults). + self.warmup = 3 + self.repeat = 10 + self.stream_delay_micro_secs = 5000 + + def budget(self): + return (self.warmup, self.repeat, self.stream_delay_micro_secs) + + @classmethod + def get(cls): + if cls.instance is None: + cls.instance = cls() + return cls.instance + + +@pytest.fixture +def fake_flashinfer(monkeypatch): + """A stand-in ``flashinfer.autotuner`` that records every ``autotune`` context.""" + _FakeTuner.instance = None + events: list[tuple] = [] + active: list[tuple] = [] + + @contextmanager + def autotune(tune_mode=True, cache=None, tuning_buckets=None, round_up=None, skip_ops=None): + entry = (tune_mode, tuning_buckets) + events.append(("enter", *entry)) + active.append(entry) + try: + yield + finally: + active.pop() + events.append(("exit", *entry)) + + module = types.ModuleType("flashinfer.autotuner") + module.autotune = autotune + module.AutoTuner = _FakeTuner + package = types.ModuleType("flashinfer") + package.autotuner = module + monkeypatch.setitem(sys.modules, "flashinfer", package) + monkeypatch.setitem(sys.modules, "flashinfer.autotuner", module) + monkeypatch.delenv(FLASHINFER_AUTOTUNE_ENV, raising=False) + monkeypatch.delenv(FLASHINFER_AUTOTUNE_BUDGET_ENV, raising=False) + # Healthy CUDA context unless a test says otherwise; no device tag (other test modules + # leave a stub ``torch`` in sys.modules, so the tag must not depend on the import). + monkeypatch.setattr(helper, "cuda_context_error", lambda device=None: None) + monkeypatch.setattr(helper, "_flashinfer_device_tag", lambda: "") + return types.SimpleNamespace(events=events, active=active, tuner=_FakeTuner) + + +def test_tunes_once_at_the_exact_token_count_then_times_under_lookup_only(fake_flashinfer): + seen: list[tuple] = [] + + # A key tuned by an earlier case on this worker is not reported again. + fake_flashinfer.tuner.get().profiling_cache[("moe", 64)] = ((64, 1), None) + + def tune_fn(): + # The tuning forward runs inside tune_mode=True with the exact bucket and + # stands in for the autotuner filling its cache. + seen.append((tuple(fake_flashinfer.active), fake_flashinfer.tuner.get().budget())) + fake_flashinfer.tuner.get().profiling_cache[("moe", 320)] = ("tactic", "profile") + + with flashinfer_autotuned(320, tune_fn, label="vllm_trtllm_nvfp4") as tuned: + assert tuned is True + assert fake_flashinfer.active == [(False, (320,))] + # The reduced budget applies to the tuning pass only; the timed body sees + # FlashInfer's own values again. + assert fake_flashinfer.tuner.get().budget() == (3, 10, 5000) + + # Tuned under the collector's default budget: 1 warm-up, 3 timed, 1 ms delay. + assert seen == [(((True, (320,)),), (1, 3, 1000))] + assert fake_flashinfer.tuner.get().budget() == (3, 10, 5000) + assert fake_flashinfer.events == [ + ("enter", True, (320,)), + ("exit", True, (320,)), + ("enter", False, (320,)), + ("exit", False, (320,)), + ] + + +def test_budget_env_overrides_or_keeps_flashinfer_defaults(fake_flashinfer, monkeypatch): + tuner = fake_flashinfer.tuner.get() + seen = [] + + monkeypatch.setenv(FLASHINFER_AUTOTUNE_BUDGET_ENV, "2, 5, 2000") + with flashinfer_autotuned(64, lambda: seen.append(tuner.budget())): + pass + assert seen == [(2, 5, 2000)] + assert tuner.budget() == (3, 10, 5000) + + monkeypatch.setenv(FLASHINFER_AUTOTUNE_BUDGET_ENV, "flashinfer") + with flashinfer_autotuned(64, lambda: seen.append(tuner.budget())): + pass + assert seen[-1] == (3, 10, 5000) + + +@pytest.mark.parametrize("raw", ["3,10", "1,0,1000"]) +def test_budget_env_rejects_malformed_values(monkeypatch, raw): + monkeypatch.setenv(FLASHINFER_AUTOTUNE_BUDGET_ENV, raw) + with pytest.raises(ValueError, match=FLASHINFER_AUTOTUNE_BUDGET_ENV): + helper.flashinfer_autotune_budget() + + +def test_env_switch_restores_the_untuned_measurement(fake_flashinfer, monkeypatch): + monkeypatch.setenv(FLASHINFER_AUTOTUNE_ENV, "0") + calls = [] + + with flashinfer_autotuned(1024, lambda: calls.append("tuned")) as tuned: + assert tuned is False + assert fake_flashinfer.active == [] + + assert calls == [] + assert fake_flashinfer.events == [] + + +def test_default_is_on(monkeypatch): + # Off by default would silently time FlashInfer's untuned fallback tactic. + monkeypatch.delenv(FLASHINFER_AUTOTUNE_ENV, raising=False) + assert helper.flashinfer_autotune_enabled() is True + + +def test_both_moe_collectors_wrap_their_timing_in_the_tuning_context(): + """The call sites are GPU-only; pin the wiring statically so a refactor cannot + silently drop the tuning pass from either collector.""" + from pathlib import Path + + collector_root = Path(helper.__file__).resolve().parent + for backend in ("vllm", "sglang"): + source = (collector_root / backend / "collect_moe.py").read_text(encoding="utf-8") + assert "flashinfer_autotuned," in source, backend + assert source.count("flashinfer_autotuned(num_tokens, ") == 1, backend + tuned_at = source.index("flashinfer_autotuned(num_tokens, ") + timed_at = source.index("benchmark_with_power(", tuned_at) + assert 0 < timed_at - tuned_at < 400, f"{backend}: tuning context must directly wrap the timing" + + +def test_poisoned_context_after_the_tuning_pass_is_reported_as_the_autotune_failure(fake_flashinfer, monkeypatch): + """gb200 job 512476645 (2026-10-27): a Llama-4 (topk=1) nvfp4 case faulted inside a + candidate tactic; FlashInfer swallowed it, the chosen tactic then failed with an opaque + 'Got CUDA error', and the next case on the worker inherited the dead context.""" + tuner = fake_flashinfer.tuner.get() + tuner.stats = types.SimpleNamespace(failed_tactics={"flashinfer::trtllm_fp4_block_scale_moe": {(16, 3), (32, 7)}}) + dead = "AcceleratorError: CUDA error: an illegal memory access was encountered" + monkeypatch.setattr(helper, "cuda_context_error", lambda device=None: dead) + timed = [] + + with pytest.raises(RuntimeError) as info, flashinfer_autotuned(1, lambda: None, label="llama4_nvfp4"): + timed.append("timed") + + msg = str(info.value) + assert "FlashInfer autotune pass at num_tokens=1 llama4_nvfp4 left the CUDA context in error state" in msg + assert "illegal memory access" in msg + assert "2 failed tactics" in msg and "(16, 3)" in msg + assert timed == [] # the measurement never ran on the dead context + assert fake_flashinfer.events == [("enter", True, (1,)), ("exit", True, (1,))] + + +def test_tune_fn_error_on_a_healthy_context_propagates_unchanged_and_restores_the_budget(fake_flashinfer): + tuner = fake_flashinfer.tuner.get() + + def tune_fn(): + raise ValueError("shape rejected") + + with pytest.raises(ValueError, match="shape rejected"), flashinfer_autotuned(8, tune_fn): + pass + assert tuner.budget() == (3, 10, 5000) + assert fake_flashinfer.active == [] + + +def test_cuda_context_error_probe_reports_a_dead_context(monkeypatch): + class _Cuda: + @staticmethod + def current_device(): + return 0 + + @staticmethod + def synchronize(device): + raise RuntimeError("CUDA error: an illegal memory access was encountered") + + fake_torch = types.SimpleNamespace(cuda=_Cuda, ones=lambda *a, **k: None) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + assert helper.cuda_context_error() == "RuntimeError: CUDA error: an illegal memory access was encountered" + + class _HealthyCuda(_Cuda): + @staticmethod + def synchronize(device): + return None + + class _Tensor: + def sum(self): + return self + + def item(self): + return 1.0 + + fake_torch = types.SimpleNamespace(cuda=_HealthyCuda, ones=lambda *a, **k: _Tensor()) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + assert helper.cuda_context_error(0) is None diff --git a/python/aisimulate/tests/unit/collector/test_helper_phase_clock.py b/python/aisimulate/tests/unit/collector/test_helper_phase_clock.py new file mode 100644 index 0000000000..b40cf0787e --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_helper_phase_clock.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from collector.helper import PhaseClock, render_phase_totals + +pytestmark = pytest.mark.unit + + +def test_phase_clock_accumulates_phases_and_point_totals(): + # The rendered text only feeds human-read log lines; this is a smoke test of the accounting. + clock = PhaseClock() + clock.mark("build") + clock.mark("build") # a second slice accumulates into the same phase + clock.absorb({"warmup": 0.5, "timed": 0.25}, prefix="bench_") + clock.absorb(None) + assert clock.phases["bench_warmup"] == 0.5 and clock.phases["bench_timed"] == 0.25 + assert clock.total() >= clock.phases["build"] + totals: dict[str, float] = {} + clock.add_to(totals) + clock.add_to(totals) + assert totals["points"] == 2 and totals["bench_timed"] == 0.5 + assert isinstance(clock.render(), str) and isinstance(render_phase_totals(totals), str) + + +def _fake_torch(monkeypatch, allocated): + import sys + import types + + state = {"allocated": allocated} + fake = types.ModuleType("torch") + fake.cuda = types.SimpleNamespace( + is_available=lambda: True, + memory_allocated=lambda: state["allocated"], + ) + monkeypatch.setitem(sys.modules, "torch", fake) + return state + + +def test_periodic_gc_collects_every_n_calls(monkeypatch): + import gc + + from collector.helper import PeriodicGarbageCollector + + _fake_torch(monkeypatch, allocated=1000) + calls = [] + monkeypatch.setattr(gc, "collect", lambda: calls.append(1) or 0) + pgc = PeriodicGarbageCollector(every=4, growth_bytes=10**9) + results = [pgc.maybe_collect() for _ in range(9)] + assert results == [False, False, False, True, False, False, False, True, False] + assert len(calls) == 2 + + +def test_periodic_gc_collects_early_on_memory_growth(monkeypatch): + import gc + + from collector.helper import PeriodicGarbageCollector + + state = _fake_torch(monkeypatch, allocated=1000) + calls = [] + monkeypatch.setattr(gc, "collect", lambda: calls.append(1) or 0) + pgc = PeriodicGarbageCollector(every=1000, growth_bytes=500) + assert pgc.maybe_collect() is False # first call sets the floor + state["allocated"] = 1400 + assert pgc.maybe_collect() is False # within the growth budget + state["allocated"] = 1600 + assert pgc.maybe_collect() is True # grew past floor + growth_bytes + assert len(calls) == 1 + assert pgc.calls == 0 + assert pgc.floor == 1000 # the floor never rises diff --git a/python/aisimulate/tests/unit/collector/test_measurement_subprocess.py b/python/aisimulate/tests/unit/collector/test_measurement_subprocess.py new file mode 100644 index 0000000000..b83b75470f --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_measurement_subprocess.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""helper.run_measurement_subprocess: a child that prints progress keeps its worker alive; +a silent one does not, and a timeout keeps the output for diagnosis.""" + +import os +import subprocess +import sys + +import pytest +from collector import helper + +pytestmark = pytest.mark.unit + + +@pytest.fixture +def beats(monkeypatch): + calls = [] + monkeypatch.setattr(helper, "_WORKER_HEARTBEAT", lambda: calls.append(1)) + return calls + + +def _child(code): + return [sys.executable, "-c", code] + + +def test_progress_beats_and_output_is_returned(beats): + code = "import time\nfor i in range(3):\n print('point', i)\n time.sleep(0.3)\n" + rc, out = helper.run_measurement_subprocess(_child(code), env=dict(os.environ), cwd=None, poll_sec=0.2) + assert rc == 0 + assert out.decode().split() == ["point", "0", "point", "1", "point", "2"] + assert beats # unbuffered prints grew the output between polls + + +def test_a_silent_child_does_not_beat(beats): + rc, _ = helper.run_measurement_subprocess( + _child("import time; time.sleep(0.8)"), env=dict(os.environ), cwd=None, poll_sec=0.2 + ) + assert rc == 0 + assert beats == [] + + +def test_timeout_kills_the_child_and_keeps_its_output(beats): + code = "import time\nprint('started')\ntime.sleep(30)\n" + with pytest.raises(subprocess.TimeoutExpired) as caught: + helper.run_measurement_subprocess(_child(code), env=dict(os.environ), cwd=None, timeout=0.6, poll_sec=0.2) + assert b"started" in caught.value.output diff --git a/python/aisimulate/tests/unit/collector/test_mhc_module_glm5next.py b/python/aisimulate/tests/unit/collector/test_mhc_module_glm5next.py index eb0d9fe77e..c2e6b7e098 100644 --- a/python/aisimulate/tests/unit/collector/test_mhc_module_glm5next.py +++ b/python/aisimulate/tests/unit/collector/test_mhc_module_glm5next.py @@ -103,18 +103,14 @@ def test_glm_row_conventions_match_the_shared_table(monkeypatch): @pytest.mark.parametrize( ("backend", "architecture", "version", "audited"), [ - ("vllm", DSV4, "0.24.0", True), - ("vllm", DSV4, "0.25.0", True), + # DeepSeek-V4 is not audited on 0.31.0 (shared architecture-free key, see the module) ("vllm", DSV4, "0.31.0", False), ("vllm", GLM, "0.31.0", True), ("vllm", GLM, "0.30.0", False), - ("vllm", GLM, "0.25.0", False), - ("vllm", GLM, "0.27.1", False), - ("sglang", DSV4, "0.5.14", True), + ("sglang", DSV4, "0.5.21", True), ("sglang", DSV4, "0.5.20", False), - ("sglang", GLM, "0.5.20", True), - ("sglang", GLM, "0.5.14", False), - ("sglang", "UnknownForCausalLM", "0.5.20", False), + ("sglang", GLM, "0.5.21", False), # not ported to 0.5.21 yet + ("sglang", "UnknownForCausalLM", "0.5.21", False), ], ) def test_each_architecture_runs_only_on_its_audited_release(monkeypatch, backend, architecture, version, audited): @@ -139,21 +135,12 @@ def test_audited_releases_never_overlap_across_architectures(monkeypatch, backen module = _load(monkeypatch, backend) from collector.version_resolver import _check_compat - candidates = ["0.24.0", "0.25.0", "0.27.1", "0.30.0", "0.31.0", "0.5.14", "0.5.16", "0.5.20"] + candidates = ["0.24.0", "0.25.0", "0.27.1", "0.30.0", "0.31.0", "0.5.14", "0.5.16", "0.5.20", "0.5.21"] for version in candidates: admitted = [arch for arch, spec in module._ARCHITECTURE_COMPAT.items() if _check_compat(spec, version)] assert len(admitted) <= 1, (version, admitted) -def test_vllm_file_compat_admits_glm_runtime_and_keeps_dsv4_releases(monkeypatch): - module = _load(monkeypatch, "vllm") - from collector.version_resolver import _check_compat - - for version in ("0.24.0", "0.25.0", "0.31.0"): - assert _check_compat(module.__compat__, version), version - assert not _check_compat(module.__compat__, "0.32.0") - - def test_vllm_glm_sites_import_the_0_31_0_module_location(): # vllm.models.glm5next.nvidia.model moved to .common.model at 0.31.0. source = (COLLECTOR_DIR / "vllm" / "collect_mhc_module.py").read_text(encoding="utf-8") diff --git a/python/aisimulate/tests/unit/collector/test_model_cases.py b/python/aisimulate/tests/unit/collector/test_model_cases.py index fd81adb5d4..82abd4897e 100644 --- a/python/aisimulate/tests/unit/collector/test_model_cases.py +++ b/python/aisimulate/tests/unit/collector/test_model_cases.py @@ -522,6 +522,7 @@ def test_moe_model_quantization_policy_is_yaml_backed(): # -BF16 suffix (there is no zai-org/GLM-5.3-FP8 repository). assert moe_model_allows_quantization("sglang", "zai-org/GLM-5.3", "fp8_block") assert not moe_model_allows_quantization("sglang", "zai-org/GLM-5.3", "bfloat16") + assert not moe_model_allows_quantization("sglang", "zai-org/GLM-5.3", "nvfp4") assert moe_model_allows_quantization("sglang", "zai-org/GLM-5.3-BF16", "bfloat16") assert not moe_model_allows_quantization("sglang", "zai-org/GLM-5.3-BF16", "fp8_block") @@ -1328,8 +1329,6 @@ def test_mla_module_metadata_and_micro_sweeps_are_yaml_backed(): ("bfloat16", "fp8", "bfloat16"), ("fp8", "fp8", "bfloat16"), ("bfloat16", "bfloat16", "fp8_block"), - ("bfloat16", "bfloat16", "fp8"), - ("bfloat16", "fp8", "fp8"), ("bfloat16", "fp8", "fp8_block"), ("fp8", "fp8", "fp8_block"), ("bfloat16", "bfloat16", "nvfp4"), @@ -1359,8 +1358,6 @@ def test_mla_module_metadata_and_micro_sweeps_are_yaml_backed(): ("bfloat16", "bfloat16", "bfloat16"), ("bfloat16", "fp8", "bfloat16"), ("bfloat16", "bfloat16", "fp8_block"), - ("bfloat16", "bfloat16", "fp8"), - ("bfloat16", "fp8", "fp8"), ("bfloat16", "fp8", "fp8_block"), # MSA-scoped combos (attention_types [msa]) at SM90: both KV dtypes # for the bf16 and fp8_block gemm tiers (fp8-KV has no SM floor for @@ -1382,8 +1379,6 @@ def test_mla_module_metadata_and_micro_sweeps_are_yaml_backed(): ("bfloat16", "bfloat16", "bfloat16"), ("bfloat16", "fp8", "bfloat16"), ("bfloat16", "bfloat16", "fp8_block"), - ("bfloat16", "bfloat16", "fp8"), - ("bfloat16", "fp8", "fp8"), ("bfloat16", "fp8", "fp8_block"), ("bfloat16", "bfloat16", "nvfp4"), ("bfloat16", "fp8", "nvfp4"), @@ -1399,8 +1394,6 @@ def test_mla_module_metadata_and_micro_sweeps_are_yaml_backed(): ("bfloat16", "fp8", "bfloat16"), ("fp8", "fp8", "bfloat16"), ("bfloat16", "bfloat16", "fp8_block"), - ("bfloat16", "bfloat16", "fp8"), - ("bfloat16", "fp8", "fp8"), ("bfloat16", "fp8", "fp8_block"), ("fp8", "fp8", "fp8_block"), ("bfloat16", "bfloat16", "nvfp4"), @@ -1422,8 +1415,8 @@ def test_mla_module_metadata_and_micro_sweeps_are_yaml_backed(): "zai-org/GLM-5.2", "zai-org/GLM-5.2-FP8", "nvidia/GLM-5.2-NVFP4", - "zai-org/GLM-5.3", "zai-org/GLM-5.3-BF16", + "zai-org/GLM-5.3", "nvidia/GLM-5.3-NVFP4", } assert {spec.native_num_heads for spec in dsa_specs if spec.architecture == "GlmMoeDsaForCausalLM"} == {64} @@ -1585,11 +1578,10 @@ def test_model_cases_path_can_infer_model_path(): assert "dsv4_csa_context_module" in plan.selected_ops assert "dsv4_csa_topk_calib" in plan.selected_ops assert "mhc_module" in plan.selected_ops - assert { - "dsv4_paged_mqa_logits_module", - "dsv4_hca_attn_module", - "dsv4_csa_attn_module", - }.isdisjoint(plan.selected_ops) + # the SDK's only live sparse sidecar (CP chunk walk) is planned; the FMLA sidecars + # (hca_attn / csa_attn) were removed from the registry on 2026-10-15 + assert "dsv4_paged_mqa_logits_module" in plan.selected_ops + assert {"dsv4_hca_attn_module", "dsv4_csa_attn_module"}.isdisjoint(plan.selected_ops) def test_plan_rejects_model_declared_ops_unknown_to_backend_registry(tmp_path): @@ -1643,7 +1635,7 @@ def test_vllm_024_schedules_consumed_dsv4_modules_only(): "dsv4_csa_generation_module", "dsv4_hca_generation_module", } - registry_only_ops = {"dsv4_paged_mqa_logits_module", "dsv4_hca_attn_module", "mhc_module"} + removed_sidecar_ops = {"dsv4_hca_attn_module", "dsv4_csa_attn_module"} # removed 2026-10-15 plan = build_collection_case_plan(backend="vllm", model_path="sgl-project/DeepSeek-V4-Pro-FP8") assert plan.ops == [ @@ -1651,12 +1643,15 @@ def test_vllm_024_schedules_consumed_dsv4_modules_only(): "dsv4_csa_generation_module", "dsv4_hca_context_module", "dsv4_hca_generation_module", + "dsv4_paged_mqa_logits_module", # consumed by the SDK CP chunk walk; shipped for vllm 0.24/0.25 "gemm", + "mhc_module", # SDK queries mhc pre/post on every backend; scheduled again 2026-10-11 "moe", ] + registered_ops = {entry.op for entry in REGISTRY} assert consumed_dsv4_ops <= plan.selected_ops - assert registry_only_ops.isdisjoint(plan.selected_ops) - assert consumed_dsv4_ops | registry_only_ops <= {entry.op for entry in REGISTRY} + assert consumed_dsv4_ops <= registered_ops + assert removed_sidecar_ops.isdisjoint(registered_ops) def test_model_architecture_can_select_case_file(): @@ -1727,8 +1722,7 @@ def test_vllm_024_model_plans_only_schedule_representable_attention_paths(): "mla_context", "mla_generation", "moe", - "moe_ep", - ] + ] # moe_ep left the default trtllm plan with the separate WideEP pin (2026-10-04) assert build_collection_case_plan(backend="vllm_xpu", model_path=kimi_path).ops == ["gemm", "moe"] models_with_unrepresentable_vllm_attention = ( @@ -1848,6 +1842,7 @@ def get_model_specs(**kwargs): "get_context_test_cases": lambda _attention_type: [[128, 1, 8, "bfloat16", "bfloat16", "bfloat16"]], "get_generation_test_cases": lambda _attention_type: [], "get_mla_module_model_specs": get_model_specs, + "get_sm_version": lambda: 90, } exec(compile(ast.Module(body=[function], type_ignores=[]), str(source_path), "exec"), namespace) @@ -1857,6 +1852,183 @@ def get_model_specs(**kwargs): assert cases == [[128, 1, 8, "bfloat16", "bfloat16", "bfloat16", "deepseek-ai/DeepSeek-V3", "mla"]] +def _load_module_case_builder(backend: str, namespace: dict): + from types import SimpleNamespace + + source_path = REPO_ROOT / f"collector/{backend}/collect_mla_module.py" + tree = ast.parse(source_path.read_text(), filename=str(source_path)) + wanted = {"_build_module_test_cases"} + if backend == "vllm": + # the DSA context grouping expands each series through the plan predicate + wanted |= {"_dsa_context_point_allowed", "_dsa_context_budgets"} + namespace.setdefault( + "get_mla_module_sweep_spec", + lambda _backend=None: SimpleNamespace( + context_dsa_chunk_prefill_size=16384, + context_max_tokens=131072, + context_dsa_max_full_sequence_length=1048576, + ), + ) + functions = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in wanted] + exec(compile(ast.Module(body=functions, type_ignores=[]), str(source_path), "exec"), namespace) + return namespace["_build_module_test_cases"] + + +@pytest.mark.parametrize("backend", ["vllm", "trtllm"]) +def test_dsa_sub_8_head_cells_stay_queued_and_raise_before_the_module_build(backend): + """b200 2026-10-06 (jobs 471043976 / 469988228): per-rank heads 1/2/4 have no sparse-MLA + kernel on SM100/103 and failed 68,435 + 15,408 cases after a full module load each. The + plan keeps them (a kernel limit is never a plan-time drop); the worker raises a + kernel_limit failure before building anything, on SM100/103 and for DSA only.""" + from contextlib import nullcontext + from types import SimpleNamespace + + heads = [128, 64, 8, 4, 2, 1] + base = [[1, 1, h, "bfloat16", "bfloat16", "bfloat16", 0] for h in heads] + ns = { + "get_context_test_cases": lambda _t: base, + "get_generation_test_cases": lambda _t: [], + "get_mla_module_model_specs": lambda **_k: [SimpleNamespace(model_path="deepseek-ai/DeepSeek-V3.2")], + "get_sm_version": lambda: 100, + "_model_max_position_embeddings": lambda _m: None, + } + build = _load_module_case_builder(backend, ns) + assert sorted({c[2] for c in build("dsa", "context")}) == sorted(heads) + + source_path = REPO_ROOT / "collector" / backend / "collect_mla_module.py" + tree = ast.parse(source_path.read_text(encoding="utf-8"), filename=str(source_path)) + worker = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "run_mla_module_worker"] + built = [] + for sm, attn_type, num_heads, raises in ( + (100, "dsa", 4, True), + (103, "dsa", 1, True), + (100, "dsa", 8, False), + (100, "mla", 4, False), + (90, "dsa", 4, False), + ): + namespace = { + "get_sm_version": lambda sm=sm: sm, + "run_mla_module": lambda **kwargs: built.append(kwargs), + "torch": SimpleNamespace(inference_mode=nullcontext), + } + exec(compile(ast.Module(body=worker, type_ignores=[]), str(source_path), "exec"), namespace) + args = (16, 1, num_heads, "bfloat16", "bfloat16", "bfloat16", "deepseek-ai/DeepSeek-V3.2", attn_type) + before = len(built) + if raises: + with pytest.raises(RuntimeError, match=r"^FIXME\(kernel-limit\): "): + namespace["run_mla_module_worker"](*args, perf_filename="x") + assert len(built) == before, (backend, sm, attn_type, num_heads) + else: + namespace["run_mla_module_worker"](*args, perf_filename="x") + assert len(built) == before + 1, (backend, sm, attn_type, num_heads) + + +@pytest.mark.parametrize("backend", ["vllm", "trtllm"]) +def test_module_plan_keeps_only_the_heads_of_declared_tensor_parallel_shards(backend): + """A model row that declares tensor_parallel_sizes (case_authoring model-correlated + narrowing, as the MSA rows do) measures only native_num_heads // tp; the shared head + grid's other entries (96/48/24/12 for the 96-head Kimi-K3 MLA module, 128 for a 64-head + model) were 40-50% of the b200 vllm dsa_context rows (2026-10-29). Rows without the + declaration keep the full grid.""" + from types import SimpleNamespace + + heads = [128, 96, 64, 48, 32, 24, 16, 12, 8] + base = [[1, 1, h, "bfloat16", "bfloat16", "bfloat16", 0] for h in heads] + glm = SimpleNamespace(model_path="zai-org/GLM-5.2", native_num_heads=64, tensor_parallel_sizes=(1, 2, 4, 8)) + kimi = SimpleNamespace(model_path="moonshotai/Kimi-K3", native_num_heads=96, tensor_parallel_sizes=None) + ns = { + "get_context_test_cases": lambda _t: base, + "get_generation_test_cases": lambda _t: [], + "get_mla_module_model_specs": lambda **_k: [glm, kimi], + "get_sm_version": lambda: 100, + "_model_max_position_embeddings": lambda _m: None, + } + build = _load_module_case_builder(backend, ns) + by_model = {} + for case in build("dsa", "context"): + by_model.setdefault(case[6], []).append(case[2]) + assert by_model["zai-org/GLM-5.2"] == [64, 32, 16, 8] + assert by_model["moonshotai/Kimi-K3"] == heads + + +def test_vllm_module_plan_groups_the_seq_len_points_of_one_module_into_one_case(): + """One case per (batch, heads, dtypes, gemm, model, attn, prefix) with its seq_len points + in slot 0, largest first: run_mla_module builds the module once per case (the module + build dominated vllm dsa_context_module in the 2026-10 campaign).""" + from types import SimpleNamespace + + base = [ + [128, 1, 64, "bfloat16", "bfloat16", "bfloat16", 0], + [16, 1, 64, "bfloat16", "bfloat16", "bfloat16", 0], + [1024, 1, 64, "bfloat16", "bfloat16", "bfloat16", 0], + [128, 1, 64, "bfloat16", "bfloat16", "bfloat16", 4096], # other prefix -> same module series + [16, 1, 64, "bfloat16", "bfloat16", "bfloat16", 4096], + [1024, 1, 64, "bfloat16", "bfloat16", "bfloat16", 4096], + [16, 8, 64, "bfloat16", "bfloat16", "bfloat16", 0], # other batch -> other case + [16, 1, 64, "fp8", "bfloat16", "bfloat16", 0], # other kv dtype -> other case + ] + ns = { + "get_context_test_cases": lambda _t: base, + "get_generation_test_cases": lambda _t: [], + "get_mla_module_model_specs": lambda **_k: [SimpleNamespace(model_path="deepseek-ai/DeepSeek-V3.2")], + "get_sm_version": lambda: 100, + "_model_max_position_embeddings": lambda _m: None, + } + build = _load_module_case_builder("vllm", ns) + cases = build("dsa", "context") + # DSA context: the module series carries its isl list AND its prefix list (the plan's + # isl x prefix grid); the runner expands the grid through the same predicate. + assert cases == [ + [[1024, 128, 16], 1, 64, "bfloat16", "bfloat16", "bfloat16", "deepseek-ai/DeepSeek-V3.2", "dsa", [4096, 0]], + [[16], 8, 64, "bfloat16", "bfloat16", "bfloat16", "deepseek-ai/DeepSeek-V3.2", "dsa", [0]], + [[16], 1, 64, "fp8", "bfloat16", "bfloat16", "deepseek-ai/DeepSeek-V3.2", "dsa", [0]], + ] + # the trtllm planner keeps one case per point + trt = _load_module_case_builder("trtllm", dict(ns)) + assert [c[0] for c in trt("dsa", "context")] == [128, 16, 1024, 128, 16, 1024, 16, 16] + + +def test_vllm_dsa_module_plan_refuses_a_grid_the_runner_would_not_reproduce(): + """The plan generator and run_mla_module share _dsa_context_point_allowed; a base list whose + (isl, prefix) points are not that predicate's product cannot be expressed as a series.""" + from types import SimpleNamespace + + base = [ + [1024, 1, 64, "bfloat16", "bfloat16", "bfloat16", 0], + [16, 1, 64, "bfloat16", "bfloat16", "bfloat16", 4096], # (1024, 4096) and (16, 0) missing + ] + ns = { + "get_context_test_cases": lambda _t: base, + "get_generation_test_cases": lambda _t: [], + "get_mla_module_model_specs": lambda **_k: [SimpleNamespace(model_path="deepseek-ai/DeepSeek-V3.2")], + "get_sm_version": lambda: 100, + "_model_max_position_embeddings": lambda _m: None, + } + build = _load_module_case_builder("vllm", ns) + with pytest.raises(RuntimeError, match="expands to 4 points but the plan holds 2"): + build("dsa", "context") + + +def test_vllm_dsa_module_plan_caps_prefix_at_the_model_max_position(): + """7,050 b200 vllm DSA cases (job 471043976) asked vLLM for max_model_len > the model's + max_position_embeddings (prefix 262144+ on DeepSeek-V3.2 / GLM-5): the planner now + applies the same per-model cap the sglang module collector applies.""" + from types import SimpleNamespace + + base = [[16, 1, 64, "bfloat16", "bfloat16", "bfloat16", p] for p in (0, 4096, 163824, 163825, 262144)] + ns = { + "get_context_test_cases": lambda _t: base, + "get_generation_test_cases": lambda _t: [], + "get_mla_module_model_specs": lambda **_k: [SimpleNamespace(model_path="deepseek-ai/DeepSeek-V3.2")], + "get_sm_version": lambda: 100, + "_model_max_position_embeddings": lambda _m: 163840, + } + build = _load_module_case_builder("vllm", ns) + assert [c[-1] for c in build("dsa", "context")] == [[163824, 4096, 0]] + ns["_model_max_position_embeddings"] = lambda _m: None # unknown ceiling -> no cap + assert [c[-1] for c in build("dsa", "context")] == [[262144, 163825, 163824, 4096, 0]] + + def test_support_matrix_models_have_model_case_aliases(): case_aliases = set() for path in (REPO_ROOT / "collector" / "cases" / "models").glob("*_cases.yaml"): @@ -2228,3 +2400,16 @@ def test_nvfp4_checkpoint_targets_msa_module_specs(monkeypatch): assert all(s.model_path == "MiniMaxAI/MiniMax-M3" for s in specs), ( "alias rows must stay keyed to the canonical model path" ) + + +def test_msa_head_axis_derives_from_declared_shards(): + """MSA is one fixed geometry (MiniMax-M3 64q/4kv/4idx): the collectors sweep + native // tp for the model row's tensor_parallel_sizes, never the shared + GQA-style head grid (owner decision 2026-09-29; TP 16 is the ceiling).""" + from collector.case_generator import get_msa_head_counts + + for backend in ("vllm", "trtllm", "sglang"): + heads = get_msa_head_counts(backend) + assert heads == [64, 32, 16, 8, 4], (backend, heads) + # every derived shard keeps the index-head count a power of two + assert all((max(1, h // 16) & (max(1, h // 16) - 1)) == 0 for h in heads) diff --git a/python/aisimulate/tests/unit/collector/test_msa_getter_contract.py b/python/aisimulate/tests/unit/collector/test_msa_getter_contract.py index b96d46e5e6..6fdf2c3fcc 100644 --- a/python/aisimulate/tests/unit/collector/test_msa_getter_contract.py +++ b/python/aisimulate/tests/unit/collector/test_msa_getter_contract.py @@ -60,10 +60,16 @@ def _stub_torch(monkeypatch): ) +class _PeriodicGCStub: + def maybe_collect(self): + return False + + def _install_helper_stub(monkeypatch): _stub_module( monkeypatch, "collector.helper", + PeriodicGarbageCollector=_PeriodicGCStub, _resolve_local_model_path=lambda path: path, benchmark_with_power=_noop, get_sm_version=lambda: 100, @@ -237,6 +243,27 @@ def test_targeted_msa_getters_keep_prefixes_and_worker_tuple_compatibility(monke assert captured[1]["perf_filename"] == "generation.txt" +@pytest.mark.parametrize( + "sm_version,expected_max_tokens", + [(90, 32768), (100, 131072)], +) +def test_trtllm_context_getter_caps_only_the_triton_reference_implementation( + monkeypatch, sm_version, expected_max_tokens +): + """mla_module_trtllm.context.msa_triton_max_tokens binds where the collector + resolves implementation="triton" (off SM100/103); fmha_sm100 systems keep + the shared context budget. isl 32768 itself stays at b=1 on both.""" + monkeypatch.setenv("COLLECTOR_MODEL_PATH", MODEL_PATH) + module = _load_backend(monkeypatch, "trtllm") + monkeypatch.setattr(module, "get_sm_version", lambda: sm_version) + + cases = module.get_msa_context_module_test_cases() + shapes = {(case[0], case[1]) for case in cases} + + assert max(s * b for s, b in shapes) == expected_max_tokens + assert (32768, 1) in shapes + + def test_vllm_public_generation_getter_applies_live_memory_filter(monkeypatch, capsys): monkeypatch.setenv("COLLECTOR_MODEL_PATH", MODEL_PATH) module = _load_backend(monkeypatch, "vllm") diff --git a/python/aisimulate/tests/unit/collector/test_msa_kernel_limit_scope.py b/python/aisimulate/tests/unit/collector/test_msa_kernel_limit_scope.py new file mode 100644 index 0000000000..f0472d92d6 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_msa_kernel_limit_scope.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The M3 index-head kernel limit is the DECODE scorer's: a guard keyed on the +head count alone dropped 1808 prefill rows on the 0.30 sweep (2026-09-28) +that the 0.29 sweep had collected.""" + +import ast +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.unit + +SRC = Path(__file__).resolve().parents[3] / "collector" / "vllm" / "collect_msa_module.py" + + +def _predicate(): + """Load only the pure predicate — the module itself imports vllm.""" + tree = ast.parse(SRC.read_text()) + fn = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "msa_decode_scorer_kernel_limit") + ns = {} + exec(compile(ast.Module(body=[fn], type_ignores=[]), str(SRC), "exec"), ns) + return ns["msa_decode_scorer_kernel_limit"] + + +@pytest.mark.parametrize("idx_heads,num_heads", [(3, 48), (6, 96)]) +def test_non_power_of_two_index_heads_fail_only_where_the_decode_scorer_runs(idx_heads, num_heads): + limit = _predicate() + assert limit(idx_heads, num_heads, is_context=False, query_len=1).startswith("FIXME(kernel-limit):") # generation + assert limit(idx_heads, num_heads, is_context=True, query_len=1) # one-token context query + assert limit(idx_heads, num_heads, is_context=True, query_len=16) is None # prefill path + assert limit(idx_heads, num_heads, is_context=True, query_len=32768) is None + + +@pytest.mark.parametrize("idx_heads", [1, 4]) +def test_power_of_two_index_heads_never_limited(idx_heads): + limit = _predicate() + for is_context, q in ((False, 1), (True, 1), (True, 4096)): + assert limit(idx_heads, idx_heads * 16, is_context=is_context, query_len=q) is None + + +def test_builder_passes_the_query_length(): + text = SRC.read_text() + assert "query_len=seq_len if is_context else 1," in text + assert "msa_decode_scorer_kernel_limit(" in text diff --git a/python/aisimulate/tests/unit/collector/test_network_layout.py b/python/aisimulate/tests/unit/collector/test_network_layout.py index 50531e3d00..4f94621153 100644 --- a/python/aisimulate/tests/unit/collector/test_network_layout.py +++ b/python/aisimulate/tests/unit/collector/test_network_layout.py @@ -83,10 +83,13 @@ def test_submit_moe_a2a_default_image_matches_manifest(): def test_submit_trtllm_alltoall_default_image_matches_manifest(): - # Same sync requirement against the trtllm pin. The collector gates the - # installed tensorrt_llm version against this pin, so drift here would - # fail every job loudly. - manifest_image = _manifest_frameworks()["trtllm"]["default"]["images"]["default"] + # Same sync requirement, against the WideEP all-to-all build base + # (trtllm_a2a: the rc20 image is the immutable base the rc11 source wheel is + # built over), NOT the stock trtllm pin — stock moved to rc29 on 2026-10-02 + # while WideEP stays on its own pins. The collector gates the installed + # tensorrt_llm version against this pin, so drift here would fail every + # job loudly. + manifest_image = _manifest_frameworks()["trtllm_a2a"]["default"]["images"]["default"] assert _default_container_image(SUBMIT_TRTLLM_ALLTOALL) == manifest_image diff --git a/python/aisimulate/tests/unit/collector/test_no_eager_graph_fallback_sites.py b/python/aisimulate/tests/unit/collector/test_no_eager_graph_fallback_sites.py new file mode 100644 index 0000000000..da68bbc3b2 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_no_eager_graph_fallback_sites.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""No collector may time a kernel eagerly behind the CUDA-graph fallback. + +``benchmark_with_power(allow_graph_fail=True)`` replaces a failed capture with +the Python launch loop and ships that floor as kernel latency; no perf table +records the helper's ``used_cuda_graph`` flag, and layer_permissions.md allows +a measurement-method degradation only when it is recorded in the output row. +The sglang trtllm_mla prefill (B1) and the trtllm GDN lanes were closed on +2026-10-11, the last four sites — trtllm mamba2 and moe, wideep moe_compute, +network trtllm alltoall — on 2026-10-12. This test keeps the count at zero for +every file under ``collector/``: any ``allow_graph_fail=`` other than the literal +``False`` (a variable could carry ``True``) is an offender. The one sound exception +is a wrapper forwarding its own keyword-only ``allow_graph_fail=False`` parameter: +every call of that wrapper passes the keyword (or the False default) and is scanned +by the same walk.""" + +import ast +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.unit + +COLLECTOR_ROOT = Path(__file__).resolve().parents[3] / "collector" + + +def _forwards_its_false_default(function, value): + """``allow_graph_fail=allow_graph_fail`` inside a function whose keyword-only parameter defaults to False.""" + if function is None or not (isinstance(value, ast.Name) and value.id == "allow_graph_fail"): + return False + for arg, default in zip(function.args.kwonlyargs, function.args.kw_defaults, strict=True): + if arg.arg == "allow_graph_fail": + return isinstance(default, ast.Constant) and default.value is False + return False + + +def test_no_collector_passes_allow_graph_fail_other_than_literal_false(): + offenders = [] + for path in sorted(COLLECTOR_ROOT.rglob("*.py")): + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + enclosing = {} # node -> innermost enclosing function (ast.walk visits outer functions first) + for function in ast.walk(tree): + if isinstance(function, ast.FunctionDef | ast.AsyncFunctionDef): + for child in ast.walk(function): + enclosing[child] = function + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + for keyword in node.keywords: + if keyword.arg != "allow_graph_fail": + continue + if isinstance(keyword.value, ast.Constant) and keyword.value.value is False: + continue + if _forwards_its_false_default(enclosing.get(node), keyword.value): + continue + offenders.append(f"{path.relative_to(COLLECTOR_ROOT)}:{node.lineno}") + assert offenders == [], f"eager-timed rows would ship from: {offenders}" diff --git a/python/aisimulate/tests/unit/collector/test_parallel_run.py b/python/aisimulate/tests/unit/collector/test_parallel_run.py index 8ff0be5b7e..b287b555fd 100644 --- a/python/aisimulate/tests/unit/collector/test_parallel_run.py +++ b/python/aisimulate/tests/unit/collector/test_parallel_run.py @@ -10,6 +10,10 @@ - Sentinel balance under repeated restarts (the core bug-fix scenario) - Lost progress tick: worker dies after recording done but before its progress_value tick (the 99/100 hang, H20 2026-07-19) +- Stall watchdog: wedged workers are killed and restarted, the queue keeps + draining; only consecutive rounds without progress close the op, and the + never-dispatched tasks are not written into the failed ledger + (b200 vllm dsa_context_module, 2026-10-11) """ import json @@ -18,6 +22,8 @@ import os import signal import sys +import threading +import time from pathlib import Path from unittest.mock import MagicMock @@ -40,9 +46,9 @@ if _COLLECTOR_DIR not in sys.path: sys.path.insert(0, _COLLECTOR_DIR) +_torch = MagicMock() +_torch.AcceleratorError = type("AcceleratorError", (Exception,), {}) if "torch" not in sys.modules: - _torch = MagicMock() - _torch.AcceleratorError = type("AcceleratorError", (Exception,), {}) sys.modules["torch"] = _torch import collect as _collect_mod @@ -76,6 +82,26 @@ def _task_fn(label, behavior, device): os.kill(os.getpid(), signal.SIGABRT) elif behavior == "error": raise ValueError(f"simulated: {label}") + elif behavior == "hang": + # a wedged kernel: never returns (Event.wait is immune to the shortened + # time.sleep the parent patched before forking, and survives SIGUSR1) + threading.Event().wait() + elif behavior.startswith("hang_unless:"): + if not os.path.exists(behavior.split(":", 1)[1]): + threading.Event().wait() + elif behavior.startswith(("beat:", "beat_then_hang:")): + # a long series task: one heartbeat per measured point for N seconds + # (benchmark_with_power beats the same way); "beat_then_hang" then + # wedges like "hang", so the watchdog timer runs from its last beat + from helper import worker_heartbeat + + kind, seconds = behavior.split(":", 1) + end = time.monotonic() + float(seconds) + while time.monotonic() < end: + worker_heartbeat() + time.sleep(0.1) + if kind == "beat_then_hang": + threading.Event().wait() elif behavior == "oom": raise sys.modules["torch"].OutOfMemoryError(f"simulated: {label}") # "normal": return silently @@ -94,6 +120,7 @@ def _worker_dying_after_done( module_name="unknown", current_task_ids=None, consumed_sentinel=None, + heartbeat=None, ): """Scripted stand-in for ``collect.worker`` reproducing the 2026-07-19 H20 death window: the task SUCCEEDS and lands in ``done_tasks``, but the @@ -161,6 +188,12 @@ def _fork_mp(monkeypatch): warnings.filterwarnings("ignore", message=".*fork.*", category=DeprecationWarning) ctx = mp.get_context("fork") monkeypatch.setattr(_collect_mod, "mp", ctx) + # In a full run other test modules import collect before this one without a + # torch stub (collect.torch stays None) or swap the stub in sys.modules; the + # fork'd workers would then die in _require_torch and be restarted until the + # restart cap, which looks like a hang. Pin the stub in both places per test. + monkeypatch.setitem(sys.modules, "torch", _torch) + monkeypatch.setattr(_collect_mod, "torch", _torch) @pytest.fixture(autouse=True) @@ -274,6 +307,10 @@ def test_torch_out_of_memory_error_is_fatal(self): "CUBLAS_STATUS_EXECUTION_FAILED", "CUBLAS_STATUS_INTERNAL_ERROR", "CUBLAS_STATUS_ALLOC_FAILED", + # FlashInfer TVM-FFI launcher wrapper (error text only on stderr) + "Error in function 'run' at /workspace/csrc/fused_moe/trtllm_backend/" + "trtllm_fused_moe_routing_llama4.cu:550: Got CUDA error. See above for details.", + "CUDA error: cudaErrorIllegalAddress", ], ) def test_cuda_fatal_markers_are_fatal(self, message): @@ -564,6 +601,87 @@ def _on_alarm(signum, frame): assert len(_crash_errors(errors)) >= 1 +class TestStallWatchdog: + """STALL_TIMEOUT_SEC of zero progress kills and restarts the workers; only + STALL_MAX_ROUNDS consecutive rounds without progress close the op, and the + tasks nobody dispatched stay out of the failed ledger.""" + + @pytest.fixture(autouse=True) + def _fast_stall(self, monkeypatch): + monkeypatch.setattr(_collect_mod, "STALL_THRESHOLD", 1) + monkeypatch.setattr(_collect_mod, "STALL_DIAG_SEC", 0) + monkeypatch.setattr(_collect_mod, "STALL_TIMEOUT_SEC", 2) + + @staticmethod + def _ids(errors, error_type): + return {e["task_id"] for e in errors if e.get("error_type") == error_type} + + def test_stall_kill_restarts_workers_and_the_queue_keeps_draining(self, tmp_path, monkeypatch): + monkeypatch.setattr(_collect_mod, "STALL_MAX_ROUNDS", 3) + module_name = "stall_restart" + # both workers take a wedged task first; the six normal tasks queue behind them + tasks = _tasks([("h0", "hang"), ("h1", "hang")] + [(f"n{i}", "normal") for i in range(6)]) + + errors = _run(tasks, 2, tmp_path, module_name=module_name) + + assert self._ids(errors, "WorkerStallTimeout") == {"h0", "h1"} + assert self._ids(errors, "WorkerOrphanedTask") == set() + assert _crash_errors(errors) == [] # a watchdog kill is not a crash + assert _load_done_ids(tmp_path, module_name) == {f"n{i}" for i in range(6)} + assert _load_failed_ids(tmp_path, module_name) == {"h0", "h1"} + + def test_heartbeats_hold_the_watchdog_until_the_measuring_task_finishes(self, tmp_path, monkeypatch): + monkeypatch.setattr(_collect_mod, "STALL_MAX_ROUNDS", 3) + module_name = "stall_heartbeat" + # "beat" measures for longer than STALL_TIMEOUT_SEC with one heartbeat per + # point and must finish; "quiet" beats briefly, then wedges: it is killed + # once nothing else beats, STALL_TIMEOUT_SEC after the last progress + tasks = _tasks([("beat", "beat:3.5"), ("quiet", "beat_then_hang:0.5")]) + + errors = _run(tasks, 2, tmp_path, module_name=module_name) + + assert self._ids(errors, "WorkerStallTimeout") == {"quiet"} + assert self._ids(errors, "WorkerOrphanedTask") == set() + assert _crash_errors(errors) == [] + assert _load_done_ids(tmp_path, module_name) == {"beat"} + assert _load_failed_ids(tmp_path, module_name) == {"quiet"} + + def test_consecutive_stall_rounds_close_the_op_and_orphans_are_not_failed(self, tmp_path, monkeypatch): + monkeypatch.setattr(_collect_mod, "STALL_MAX_ROUNDS", 2) + module_name = "stall_close" + release = tmp_path / "release" + tasks = _tasks([(f"h{i}", f"hang_unless:{release}") for i in range(6)]) + + errors = _run(tasks, 2, tmp_path, module_name=module_name) + + # round 1: h0/h1 wedged -> killed, workers restarted; round 2: h2/h3 + # wedged with nothing finished in between -> op closed. h4/h5 were + # never dispatched: reported, but neither failed nor attempted. + stalled = self._ids(errors, "WorkerStallTimeout") + orphaned = self._ids(errors, "WorkerOrphanedTask") + assert stalled == {"h0", "h1", "h2", "h3"} + assert orphaned == {"h4", "h5"} + assert _load_done_ids(tmp_path, module_name) == set() + assert _load_failed_ids(tmp_path, module_name) == stalled + assert _load_attempted_ids(tmp_path, module_name) & orphaned == set() + + # a plain --resume (no retry-failed) runs exactly the orphans + release.touch() + resumed = parallel_run( + tasks, + _task_fn, + num_processes=2, + module_name=module_name, + resume_options={"checkpoint_dir": str(tmp_path / ".checkpoint"), "resume": True}, + ) + assert _load_done_ids(tmp_path, module_name) == {"h4", "h5"} + assert _load_failed_ids(tmp_path, module_name) == stalled + assert self._ids(resumed, "WorkerOrphanedTask") == set() + assert self._ids(resumed, "WorkerStallTimeout") == set() + unresolved = [e for e in resumed if e.get("error_type") == "UnresolvedFailures"] + assert len(unresolved) == 1 and "4 unresolved" in unresolved[0]["error_message"] + + class TestResumeIntegrity: """Checkpoint binds runtime identity; resume surfaces unresolved failures.""" @@ -665,3 +783,26 @@ def test_fresh_zero_work_run_replaces_stale_checkpoint(self, tmp_path): assert _load_done_ids(tmp_path, module_name) == set() assert _load_failed_ids(tmp_path, module_name) == set() assert _load_attempted_ids(tmp_path, module_name) == set() + + +def test_worker_heartbeat_finds_the_sink_registered_through_the_other_module_name(monkeypatch): + """collect.py registers the sink on ``helper``; the op modules beat through + ``collector.helper`` (a second module object) and must still reach it.""" + import types + + import helper + + beats = [] + twin = types.ModuleType("collector.helper") + twin._WORKER_HEARTBEAT = lambda: beats.append(1) + monkeypatch.setattr(helper, "_WORKER_HEARTBEAT", None) + monkeypatch.setitem(sys.modules, "collector.helper", twin) + + helper.worker_heartbeat() + helper.worker_heartbeat() + + assert beats == [1, 1] + monkeypatch.setattr(helper, "_WORKER_HEARTBEAT", None) + monkeypatch.delitem(sys.modules, "collector.helper") + helper.worker_heartbeat() # no sink anywhere: a silent no-op + assert beats == [1, 1] diff --git a/python/aisimulate/tests/unit/collector/test_perf_output_finalization.py b/python/aisimulate/tests/unit/collector/test_perf_output_finalization.py index 3205f270ab..9311db5887 100644 --- a/python/aisimulate/tests/unit/collector/test_perf_output_finalization.py +++ b/python/aisimulate/tests/unit/collector/test_perf_output_finalization.py @@ -880,6 +880,84 @@ def test_convert_without_merge_overwrites(tmp_path): assert _rows_by_key(parquet) == {"s2": 2.0} # legacy overwrite preserved when opted out +def test_finalize_third_merge_survives_large_string_round_trip(tmp_path): + """run -> resume -> retry-failed: the second merge writes string columns back as + `large_string` (pandas round trip); the third finalize must still merge, not + overwrite (sm120 attention_generation fp8kv lost 27,303 rows this way, 2026-10-01).""" + import pyarrow as pa + + perf = tmp_path / "gemm_perf.txt" + parquet = perf.with_suffix(".parquet") + _write_keyed_perf_csv(perf, [("s1", 1.0)]) + finalize_perf_files([perf]) + _write_keyed_perf_csv(perf, [("s2", 2.0)]) + finalize_perf_files([perf]) + # force the state the pandas round trip produces: identity column stored as large_string + t = pq.read_table(parquet) + idx = t.schema.get_field_index("shape") + pq.write_table(t.set_column(idx, "shape", t.column("shape").cast(pa.large_string())), parquet) + assert str(pq.read_table(parquet).schema.field("shape").type) == "large_string" + _write_keyed_perf_csv(perf, [("s3", 3.0)]) + finalize_perf_files([perf]) + assert _rows_by_key(parquet) == {"s1": 1.0, "s2": 2.0, "s3": 3.0} + + +@contextmanager +def _pandas_unavailable(monkeypatch): + """The collector images ship pyarrow (harness ``pyarrow-site``) but no pandas.""" + import builtins + + real_import = builtins.__import__ + + def _import(name, *args, **kwargs): + if name == "pandas" or name.startswith("pandas."): + raise ModuleNotFoundError("No module named 'pandas'") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", _import) + try: + yield + finally: + monkeypatch.setattr(builtins, "__import__", real_import) + + +def test_finalize_merge_does_not_need_pandas(tmp_path, monkeypatch): + """The second finalize of a shard (resume / retry-failed) merges into the + existing parquet; it died with ``No module named 'pandas'`` on the vllm + image (b200 dsa_context_module job 481670367 window 2, 2026-10-12).""" + perf = tmp_path / "gemm_perf.txt" + parquet = perf.with_suffix(".parquet") + _write_keyed_perf_csv(perf, [("s1", 1.0), ("s2", 2.0)]) + finalize_perf_files([perf]) + + with _pandas_unavailable(monkeypatch): + _write_keyed_perf_csv(perf, [("s1", 9.0), ("s3", 3.0)]) + finalize_perf_files([perf]) + + assert _rows_by_key(parquet) == {"s1": 9.0, "s2": 2.0, "s3": 3.0} + table = pq.read_table(parquet) + assert [r["shape"] for r in table.to_pylist()] == ["s2", "s1", "s3"] # kept rows keep their order + assert not table.schema.metadata # no pandas index metadata in the published parquet + + +def test_finalize_moe_legacy_merge_does_not_need_pandas(tmp_path, monkeypatch): + perf = tmp_path / "moe_perf.txt" + parquet = perf.with_suffix(".parquet") + perf.write_text("kernel_source,shape,latency,default_eligible\nexact,s1,1.0,false\nexact,s2,2.0,false\n") + finalize_perf_files([perf]) + + with _pandas_unavailable(monkeypatch): + perf.write_text("kernel_source,shape,latency\nexact,s1,3.0\nexact,s3,4.0\n") # legacy recollection + finalize_perf_files([perf]) + + rows = {r["shape"]: r for r in pq.read_table(parquet).to_pylist()} + assert {k: (v["latency"], v["default_eligible"]) for k, v in rows.items()} == { + "s1": (3.0, False), # re-measured, keeps its annotation + "s2": (2.0, False), + "s3": (4.0, True), # new key: eligible by default + } + + def test_finalize_merge_falls_back_to_overwrite_on_schema_mismatch(tmp_path): perf = tmp_path / "gemm_perf.txt" parquet = perf.with_suffix(".parquet") @@ -1044,3 +1122,46 @@ def test_finalize_keeps_power_disabled_schema_column_free(tmp_path): [parquet] = finalize_perf_files([perf]) assert pq.read_table(parquet).column_names == ["shape", "latency"] + + +def test_finalize_drops_rows_re_appended_by_a_retried_task_keeping_the_newest(tmp_path): + """The staging CSV is append-only across resumed windows and a task retried with + --resume-retry-failed logs its successful rows again (rtxpro msa_context: 10,342 of + 13,972 rows were such re-appends). Finalization keeps the newest row per identity + key, in the original order, like the merge-with-existing path already did.""" + perf = tmp_path / "gemm_perf.txt" + _write_keyed_perf_csv(perf, [("s1", 1.0), ("s2", 2.0), ("s1", 1.5), ("s3", 3.0), ("s2", 2.5)]) + finalization_info = {} + + [parquet] = finalize_perf_files([perf], delete_source=False, finalization_info=finalization_info) + + table = pq.read_table(parquet) + assert table.column("shape").to_pylist() == ["s1", "s3", "s2"] + assert table.column("latency").to_pylist() == [1.5, 3.0, 2.5] + [info] = finalization_info.values() + assert info.new_rows == 3 # the deduplicated count is what gets reported + + # Rows that differ only in metric columns (power, latency) share one identity: the newer one wins. + metric_only = tmp_path / "metric_only" + metric_only.mkdir() + perf = metric_only / "gemm_perf.txt" + perf.write_text("op,shape,power,latency\nmatmul,s1,100.0,1.0\nmatmul,s1,120.0,1.1\nmatmul,s2,100.0,2.0\n") + [parquet] = finalize_perf_files([perf], delete_source=False) + assert pq.read_table(parquet).column("latency").to_pylist() == [1.1, 2.0] + + +def test_finalize_dedupe_ignores_moe_eligibility_metadata_and_tolerates_empty_columns(tmp_path): + perf = tmp_path / "moe_perf.txt" + perf.write_text( + "op,shape,kernel_source,note,default_eligible,latency\n" + "moe,s1,cutlass,,true,1.0\n" + "moe,s1,cutlass,,false,1.2\n" + "moe,s2,cutlass,,true,2.0\n" + ) + + [parquet] = finalize_perf_files([perf], delete_source=False) + + table = pq.read_table(parquet) + assert table.column("shape").to_pylist() == ["s1", "s2"] + assert table.column("default_eligible").to_pylist() == [False, True] + assert table.column("latency").to_pylist() == [1.2, 2.0] diff --git a/python/aisimulate/tests/unit/collector/test_provenance.py b/python/aisimulate/tests/unit/collector/test_provenance.py index bdbe13ead0..48e248a2ce 100644 --- a/python/aisimulate/tests/unit/collector/test_provenance.py +++ b/python/aisimulate/tests/unit/collector/test_provenance.py @@ -162,7 +162,6 @@ def test_dsa_collector_hash_covers_worker_binding(tmp_path, backend, binding_fil "collect_gemm", "collect_mhc_module", "collect_mla_module", - "collect_mla_module_027", "collect_moe", "collect_msa_module", ], @@ -179,18 +178,6 @@ def test_vllm_collector_hash_covers_native_setup(tmp_path, module_name): assert provenance.collector_hash(module, tmp_path, closures) != before -def test_vllm_mla_shim_hash_covers_executed_module(tmp_path): - module = "collector.vllm.collect_mla_module_027" - closures = provenance.load_closures(HASH_CLOSURES_PATH) - paths = {module.replace(".", "/") + ".py", *provenance.SHARED_CORE, *closures[module]} - for relative in paths - {provenance.MODEL_CASES_GROUP}: - _write(tmp_path / relative, (REPO_ROOT / relative).read_text(encoding="utf-8")) - before = provenance.collector_hash(module, tmp_path, closures) - implementation = tmp_path / "collector/vllm/collect_mla_module.py" - _write(implementation, implementation.read_text(encoding="utf-8") + "\n# native forward changed\n") - assert provenance.collector_hash(module, tmp_path, closures) != before - - def test_collector_hash_changes_when_model_cases_group_changes(tmp_path): _build_fake_repo(tmp_path) before = provenance.collector_hash("collector.sglang.collect_gemm", tmp_path, FAKE_CLOSURES) diff --git a/python/aisimulate/tests/unit/collector/test_registry_table_producers.py b/python/aisimulate/tests/unit/collector/test_registry_table_producers.py new file mode 100644 index 0000000000..1c4dac5e0e --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_registry_table_producers.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A checkpoint producer owns every table its registry entry declares. + +collect.py's fail-closed finalize binds each staged ``*_perf.txt`` to the +registry entries that declare it. A collector that writes a second table +its op did not declare cannot be finalized (first full vllm compute_scale +run, 2026-09-27: ``scale_matrix_perf.txt has no selected producer``; the +PerfFile.SCALE_MATRIX enum alone never fixed that). Owner decision +2026-09-28: the executor supports multi-table producers — ``OpEntry. +extra_perf_filenames`` — instead of splitting the op. +""" + +import importlib +import sys +from pathlib import Path + +import pytest +from collector.registry_types import OpEntry, PerfFile + +pytestmark = pytest.mark.unit + +BACKENDS = ("vllm", "sglang", "trtllm") + + +def _registry(backend): + return importlib.import_module(f"collector.{backend}.registry").REGISTRY + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_compute_scale_declares_both_tables_on_one_producer(backend): + entries = {entry.op: entry for entry in _registry(backend)} + assert "scale_matrix" not in entries # one measurement, one op + cs = entries["compute_scale"] + assert cs.perf_filename == PerfFile.COMPUTESCALE + assert cs.extra_perf_filenames == (PerfFile.SCALE_MATRIX,) + assert cs.perf_filenames == (PerfFile.COMPUTESCALE, PerfFile.SCALE_MATRIX) + + +def test_op_entry_rejects_duplicate_tables(): + with pytest.raises(ValueError, match="duplicate perf tables"): + OpEntry( + op="x", + module="m", + get_func="g", + run_func="r", + perf_filename=PerfFile.GEMM, + extra_perf_filenames=(PerfFile.GEMM,), + ) + + +def test_provenance_collections_carry_extra_tables(): + from collector.version_resolver import build_collections + + (cs,) = [ + c + for c in build_collections(_registry("vllm"), "vllm", "0.30.0", ops=["compute_scale"]) + if c["type"] == "compute_scale" + ] + assert cs["perf_filename"] == PerfFile.COMPUTESCALE + assert tuple(cs["extra_perf_filenames"]) == (PerfFile.SCALE_MATRIX,) + + +def test_executor_binds_every_declared_table_to_the_producer(): + if "torch" not in sys.modules: # collect.py imports torch at module level (see test_collect_provenance_writer) + from unittest.mock import MagicMock + + _torch = MagicMock() + _torch.AcceleratorError = type("AcceleratorError", (Exception,), {}) + sys.modules["torch"] = _torch + collector_dir = str(Path(__file__).resolve().parents[3] / "collector") + if collector_dir not in sys.path: + sys.path.insert(0, collector_dir) + import collect + + collection = { + "name": "vllm", + "type": "compute_scale", + "module": "collector.vllm.collect_computescale", + "get_func": "get_computescale_test_cases", + "run_func": "run_computescale", + "perf_filename": PerfFile.COMPUTESCALE, + "extra_perf_filenames": (PerfFile.SCALE_MATRIX,), + } + assert collect._collection_perf_filenames(collection) == ["computescale_perf.txt", "scale_matrix_perf.txt"] + assert collect._collection_tables(collection) == ("computescale_perf", "scale_matrix_perf") + identity = { + "schema": collect.RESUME_SCHEMA_VERSION, + "backend": "vllm", + "module": "vllm.compute_scale", + "run_func": "run_computescale", + "framework_version": "0.30.0", + "sm_version": 90, + } + assert set(identity) == set(collect._CHECKPOINT_IDENTITY_FIELDS) + assert collect._registered_checkpoint_tables(identity, backend="vllm") == {"computescale_perf", "scale_matrix_perf"} + gemm_identity = {**identity, "module": "vllm.gemm", "run_func": "run_gemm"} + assert collect._registered_checkpoint_tables(gemm_identity, backend="vllm") == {"gemm_perf"} + with pytest.raises(RuntimeError, match="no unambiguous registered table set"): + collect._registered_checkpoint_tables({**identity, "module": "vllm.nope"}, backend="vllm") + + +# ── shared-table producers: sglang dsa_*_module_skip_indexer ───────────────── +# The skip-indexer rows live in the full-module table (op_name "_skip_indexer" +# suffix; the SDK reads them from there), so the skip ops declare that table +# and select the mode with OpEntry.run_kwargs. Until 2026-10-06 they declared a +# dsa_*_skip_indexer_perf.txt nobody wrote and --resume finalization raised +# "open checkpoint event has no regular staging table" on every skip shard +# (b200_sxm sglang 0.5.21, GitLab job 469988017). + + +def test_sglang_skip_indexer_ops_share_the_full_module_table_and_select_by_run_kwargs(): + entries = {entry.op: entry for entry in _registry("sglang")} + for phase in ("context", "generation"): + full = entries[f"dsa_{phase}_module"] + skip = entries[f"dsa_{phase}_module_skip_indexer"] + assert skip.perf_filename == full.perf_filename == getattr(PerfFile, f"DSA_{phase.upper()}_MODULE") + assert skip.run_func == full.run_func == "run_mla_module_worker" + assert skip.run_kwargs_dict == {"skip_indexer": True} + assert full.run_kwargs == () + assert not any("SKIP_INDEXER" in member.name for member in PerfFile) + + +def test_provenance_collections_carry_run_kwargs(): + from collector.version_resolver import build_collections + + by_type = { + c["type"]: c + for c in build_collections( + _registry("sglang"), + "sglang", + "0.5.21", + ops=["dsa_context_module", "dsa_context_module_skip_indexer"], + ) + } + assert dict(by_type["dsa_context_module_skip_indexer"]["run_kwargs"]) == {"skip_indexer": True} + assert tuple(by_type["dsa_context_module"]["run_kwargs"]) == () + assert by_type["dsa_context_module_skip_indexer"]["perf_filename"] == by_type["dsa_context_module"]["perf_filename"] + + +@pytest.mark.parametrize("bad", [(("perf_filename", "x.txt"),), (("a", 1), ("a", 2))]) +def test_op_entry_rejects_executor_owned_or_malformed_run_kwargs(bad): + with pytest.raises(ValueError, match="run_kwargs"): + OpEntry(op="x", module="m", get_func="g", run_func="r", perf_filename=PerfFile.GEMM, run_kwargs=bad) diff --git a/python/aisimulate/tests/unit/collector/test_resume_finalize_zero_passed.py b/python/aisimulate/tests/unit/collector/test_resume_finalize_zero_passed.py new file mode 100644 index 0000000000..545ffbd70d --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_resume_finalize_zero_passed.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""--resume finalization of a producer whose every attempted case failed. + +b200_sxm sglang 0.5.21 (GitLab job 469988017, 2026-10-05): moe/int4_wo had +3,078 attempted / 0 passed (lane guard) and therefore no staging table; +``_pending_resume_perf_outputs`` raised "open checkpoint event has no regular +staging table" and the run died in finalization although the failures were +already in the error report. An open event with completed cases and no table +is still a hard error (rows were produced and are missing). + +Runs with ``PYTHONPATH=collector`` like test_registry_table_producers (collect.py +is a script-style module importing torch at module level; torch is stubbed). +""" + +import sys +import types +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.unit + + +def _collect(): + if "torch" not in sys.modules: + from unittest.mock import MagicMock + + _torch = MagicMock() + _torch.AcceleratorError = type("AcceleratorError", (Exception,), {}) + sys.modules["torch"] = _torch + import collect + + return collect + + +def _fake_tracker(attempted, done): + return types.SimpleNamespace(_attempted=set(attempted), _done=set(done), _failed=set(attempted) - set(done)) + + +@pytest.fixture +def provenance_ctx(): + return { + "collections": [ + { + "name": "sglang", + "type": "moe", + "module": "collector.sglang.collect_moe", + "get_func": "g", + "run_func": "run_moe_torch", + "perf_filename": "moe_perf.txt", + "extra_perf_filenames": (), + }, + ] + } + + +def test_all_failed_producer_without_staging_table_is_skipped_not_fatal(tmp_path, monkeypatch, provenance_ctx): + collect = _collect() + monkeypatch.setattr( + collect, "_load_selected_producer_checkpoint", lambda *a, **k: _fake_tracker(attempted={"a", "b"}, done=set()) + ) + pending = collect._pending_resume_perf_outputs( + tmp_path, provenance_ctx, backend="sglang", checkpoint_dir="ckpt", sm_version=100 + ) + assert pending == [] + + +def test_completed_cases_without_staging_table_still_fail_closed(tmp_path, monkeypatch, provenance_ctx): + collect = _collect() + monkeypatch.setattr( + collect, "_load_selected_producer_checkpoint", lambda *a, **k: _fake_tracker(attempted={"a", "b"}, done={"a"}) + ) + with pytest.raises(RuntimeError, match="open checkpoint event has no regular staging table"): + collect._pending_resume_perf_outputs( + tmp_path, provenance_ctx, backend="sglang", checkpoint_dir="ckpt", sm_version=100 + ) + + +def test_present_staging_table_with_open_event_is_selected(tmp_path, monkeypatch, provenance_ctx): + collect = _collect() + staging = tmp_path / "moe_perf.txt" + staging.write_text("row\n") + monkeypatch.setattr( + collect, "_load_selected_producer_checkpoint", lambda *a, **k: _fake_tracker(attempted={"a"}, done={"a"}) + ) + pending = collect._pending_resume_perf_outputs( + tmp_path, provenance_ctx, backend="sglang", checkpoint_dir="ckpt", sm_version=100 + ) + assert pending == [Path(staging)] diff --git a/python/aisimulate/tests/unit/collector/test_sglang_attention_0514.py b/python/aisimulate/tests/unit/collector/test_sglang_attention_0514.py index 7c51b4b811..29e06ccb46 100644 --- a/python/aisimulate/tests/unit/collector/test_sglang_attention_0514.py +++ b/python/aisimulate/tests/unit/collector/test_sglang_attention_0514.py @@ -323,6 +323,7 @@ def test_sglang_attention_runtime_fields_are_not_persisted_dimensions(): @pytest.mark.unit +@pytest.mark.parametrize("kernel_source", ["triton", "flashinfer"]) # flashinfer refuses it in the worker @pytest.mark.parametrize( ("sm_version", "num_heads", "num_kv_heads", "head_dim", "max_kv_elements"), [ @@ -333,6 +334,7 @@ def test_sglang_attention_runtime_fields_are_not_persisted_dimensions(): ], ) def test_blackwell_context_cases_are_not_silently_removed( + kernel_source, sm_version, num_heads, num_kv_heads, @@ -358,7 +360,7 @@ def test_blackwell_context_cases_are_not_silently_removed( attention_chunk_size=None, has_attention_sink=False, scaling=None, - kernel_source="triton", + kernel_source=kernel_source, architecture="TestForCausalLM", ) getter = _load_collector_function( @@ -388,12 +390,12 @@ def _run_attention_torch_source() -> str: @pytest.mark.unit def test_fp8_prefill_labels_follow_backend_compute_dtype(): - """FP8 prefill labeling mirrors SGLang 0.5.14 per-backend compute truth. + """FP8 prefill labeling mirrors SGLang 0.5.21 per-backend compute truth. flashinfer has no FP8 prefill compute path at all (BF16 Q reads the FP8 KV cache with descales) and must reject the fp8_context_fmha case. trtllm_mha - (TRTLLM-GEN) requires BF16 inputs but quantizes Q internally via - scaled_fp8_quant (trtllm_mha_backend.py:154-158, 291-301), so its + (TRTLLM-GEN) takes BF16 inputs but casts Q to FP8 itself + (trtllm_mha_backend.py:1600-1609), so its fp8_context_fmha case runs WITHOUT an external cast — the FP8 compute happens inside the backend. """ @@ -404,6 +406,8 @@ def test_fp8_prefill_labels_follow_backend_compute_dtype(): # trtllm_mha must be exempt from the external cast, not rejected. assert 'attn_backend_name != "trtllm_mha"' in source assert 'attn_backend_name in {"flashinfer", "trtllm_mha"}' not in source + # the refusal runs before the model runner is built: the queued case fails cheaply + assert source.index("has no FP8 prefill compute path") < source.index("MockModelRunner(") @pytest.mark.unit @@ -411,7 +415,7 @@ def test_bf16_prefill_on_fp8_kv_fails_closed_on_internally_quantizing_backends() """A BF16-compute prefill on an FP8 KV cache does not exist on fa3/trtllm_mha. fa3 casts Q to the KV dtype whenever the cache is FP8 and head_dim <= 256 - (flashattention_backend.py:857-872); TRTLLM-GEN always quantizes Q when the + (flashattention_backend.py:1322-1332); TRTLLM-GEN always quantizes Q when the cache is FP8. Collecting that combination would record FP8 compute under a bfloat16 attn_dtype label, so it must raise instead. """ diff --git a/python/aisimulate/tests/unit/collector/test_sglang_glm5_sparse_0514.py b/python/aisimulate/tests/unit/collector/test_sglang_glm5_sparse_0514.py index 856195015d..c7c162dd99 100644 --- a/python/aisimulate/tests/unit/collector/test_sglang_glm5_sparse_0514.py +++ b/python/aisimulate/tests/unit/collector/test_sglang_glm5_sparse_0514.py @@ -58,10 +58,10 @@ def test_glm5_flash_mla_keeps_the_production_topk_width(): assert "if k % 64" not in source -def test_glm5_sparse_collector_is_pinned_to_sglang_0514_dsa_api(): +def test_glm5_sparse_collector_uses_the_dsa_fuse_topk_switch_not_the_nsa_spelling(): + # the file pin itself is enforced by test_stock_collectors_are_exactly_pinned_to_their_resolved_runtime source = SOURCE_PATH.read_text(encoding="utf-8") - assert '__compat__ = "sglang==0.5.14"' in source assert "SGLANG_DSA_FUSE_TOPK" in source assert "SGLANG_NSA_FUSE_TOPK" not in source @@ -120,9 +120,10 @@ def test_glm5_sparse_platform_gating_lives_outside_the_getter(): def test_dsv4_sparse_platform_gating_lives_outside_the_getter(): - """Same contract for the DSV4 sparse sub-kernel family; the flash_mla-less - HCA/CSA collectors are parked whole-op unverified instead of silently - enumerating zero cases.""" + """Same contract for the DSV4 sparse sub-kernel family. The flash_mla-less + HCA/CSA FMLA sidecar ops were parked unverified until 2026-10-15 and are now + removed outright (retired SDK placeholders that failed every registry-generated + shard plan), so only the indexer scoring kernel and the topk calibration remain.""" from collector.capabilities import _load_capabilities from collector.sglang.registry import REGISTRY @@ -131,11 +132,10 @@ def test_dsv4_sparse_platform_gating_lives_outside_the_getter(): assert "_dsv4_topk_kernel_supported" not in dsv4_source _, op_min_sm = _load_capabilities() - floors = ("dsv4_paged_mqa_logits_module", "dsv4_hca_attn_module", "dsv4_csa_attn_module", "dsv4_csa_topk_calib") + floors = ("dsv4_paged_mqa_logits_module", "dsv4_csa_topk_calib") for op in floors: assert op_min_sm.get(op) == 90 by_op = {entry.op: entry for entry in REGISTRY} assert by_op["dsv4_paged_mqa_logits_module"].unverified_sms == (120,) assert by_op["dsv4_csa_topk_calib"].unverified_sms == (120,) - assert by_op["dsv4_hca_attn_module"].unverified is True - assert by_op["dsv4_csa_attn_module"].unverified is True + assert not {"dsv4_hca_attn_module", "dsv4_csa_attn_module"} & set(by_op) diff --git a/python/aisimulate/tests/unit/collector/test_standards_exceptions_verified_on.py b/python/aisimulate/tests/unit/collector/test_standards_exceptions_verified_on.py new file mode 100644 index 0000000000..5f1527e241 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_standards_exceptions_verified_on.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""A standards exception scoped to ``versions: all`` is bounded only by ``verified_on`` +(layer_permissions.md "Governed exceptions" item 3): it must name the collector pin of every +table it covers, so a pin bump that nobody re-verified fails here instead of being covered silently.""" + +from pathlib import Path + +import pytest +import yaml +from collector.framework_manifest import MANIFEST_PATH, _registry_entries, _resolve_from, load_manifest +from collector.op_catalog import CATALOG_PATH, load_family_map + +pytestmark = pytest.mark.unit +EXCEPTIONS = Path(__file__).resolve().parents[3] / "collector" / "evidence_exceptions.yaml" +ENTRIES = yaml.safe_load(EXCEPTIONS.read_text())["standards_exceptions"] + + +def _collector_pins(backend: str, table: str) -> set[str]: + manifest, family_map = load_manifest(MANIFEST_PATH), load_family_map(CATALOG_PATH) + return { + _resolve_from(manifest, family_map, backend, entry).version + for entry in _registry_entries(backend) + if table in {Path(str(name)).stem for name in entry.perf_filenames} + } + + +@pytest.mark.parametrize("entry", ENTRIES, ids=[e["rule"].split("—")[-1].strip() for e in ENTRIES]) +def test_verified_on_is_the_collector_pin_of_every_table_in_scope(entry): + scope = entry["scope"] + for table in scope["tables"]: + pins = _collector_pins(scope["backend"], table) + assert pins, f"{scope['backend']} has no collector for {table}" + assert pins == {entry["verified_on"]}, ( + f"{table}: collected on {sorted(pins)}, verified on {entry['verified_on']}; re-verify the " + "deviation on the new pin and move verified_on, or delete the entry" + ) diff --git a/python/aisimulate/tests/unit/collector/test_version_resolver.py b/python/aisimulate/tests/unit/collector/test_version_resolver.py index 10c3e714e4..7aecd5a68a 100644 --- a/python/aisimulate/tests/unit/collector/test_version_resolver.py +++ b/python/aisimulate/tests/unit/collector/test_version_resolver.py @@ -143,7 +143,11 @@ def test_invalid_compat_raises(self, compat): # --------------------------------------------------------------------------- @pytest.mark.unit class TestExactVersionSet: - """``_check_compat``'s grammar is AND-of-comparators only (see + """Grammar semantics only: collector modules no longer use this form (one + exact ``==`` pin per file, layer_permissions.md "Version contract"); the + tests below keep ``_check_compat``'s range/``!=`` behaviour documented. + + ``_check_compat``'s grammar is AND-of-comparators only (see ``_parse_compat_specifier``): no OR, so a __compat__ string can never directly express a discrete set like "0.5.14 or 0.5.17". When a module's verified-compatible versions are non-contiguous -- e.g. @@ -443,32 +447,21 @@ def test_output_dict_shape(self): "get_func", "run_func", "perf_filename", - "worker_perf_filename", + "extra_perf_filenames", + "run_kwargs", "unverified", "unverified_sms", } + assert c["extra_perf_filenames"] == () + assert c["run_kwargs"] == () assert c["name"] == "vllm" assert c["unverified"] is False assert c["unverified_sms"] == () - assert c["worker_perf_filename"] is None def test_perf_filename_propagated(self): colls = build_collections(self.SAMPLE_REGISTRY, "vllm", "0.17.0", ops=["gemm"]) assert colls[0]["perf_filename"] == PerfFile.GEMM - def test_worker_filename_keeps_canonical_output_filename(self): - entry = OpEntry( - op="dsa_context_module_skip_indexer", - module="collector.sglang.collect_mla_module", - get_func="get_dsa_context_module_skip_indexer_test_cases", - run_func="run_mla_module_worker", - perf_filename=PerfFile.DSA_CONTEXT_MODULE, - worker_perf_filename=PerfFile.DSA_CONTEXT_MODULE_SKIP_INDEXER, - ) - collection = build_collections([entry], "sglang", "0.5.14")[0] - assert collection["perf_filename"] == PerfFile.DSA_CONTEXT_MODULE - assert collection["worker_perf_filename"] == PerfFile.DSA_CONTEXT_MODULE_SKIP_INDEXER - # --------------------------------------------------------------------------- # OpEntry validation diff --git a/python/aisimulate/tests/unit/collector/test_vllm_attn_kv_cache_layout.py b/python/aisimulate/tests/unit/collector/test_vllm_attn_kv_cache_layout.py new file mode 100644 index 0000000000..b4a1b4db25 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_vllm_attn_kv_cache_layout.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""collect_attn.kv_cache_in_resolved_layout: the collector's paged KV cache must be laid out +in the layout the framework RESOLVED for the backend under test (RFC #42082), not a fixed HND. + +Origin: SM120 collection 2026-10-01 — FlashInfer resolves NHD outside SM100 and its XQA +decode asserted on the collector's contiguous [B, H, N, 2D] memory for every fp8-KV GQA +generation case (20,458 failures); SM100 (HND) and FLASH_ATTN never noticed. +The helper is pure torch, so the test imports it through ast, without vLLM.""" + +import ast +from pathlib import Path + +import pytest + +try: + import torch +except ImportError: + torch = None +if not isinstance(getattr(torch, "__version__", None), str): # absent, or the collector conftest's Mock stub + # reason string must match .github/application-test-inventory.json optional_collection_skips + pytest.skip("real torch required for tensor operations", allow_module_level=True) +pytestmark = pytest.mark.unit + +_SRC = Path(__file__).resolve().parents[3] / "collector" / "vllm" / "collect_attn.py" + + +def _helper(): + tree = ast.parse(_SRC.read_text()) + fn = next(n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) and n.name == "kv_cache_in_resolved_layout") + ns = {"torch": torch} + exec(compile(ast.Module(body=[fn], type_ignores=[]), str(_SRC), "exec"), ns) + return ns["kv_cache_in_resolved_layout"] + + +# KVCacheLayout.layer_view_order values (vllm/v1/kv_cache_layout.py @v0.30.0): logical [B, H, N, C] +LBHNC = (0, 1, 2, 3) # HND, identity +LBNHC = (0, 2, 1, 3) # NHD + + +@pytest.mark.parametrize("order", [LBHNC, LBNHC, (1, 0, 2, 3)]) +def test_logical_view_is_unchanged_and_physical_order_is_contiguous(order): + f = _helper() + kv = torch.arange(2 * 3 * 4 * 6, dtype=torch.float32).reshape(2, 3, 4, 6) # [B, H, N, C] + out = f(kv, order) + assert out.shape == kv.shape and torch.equal(out, kv) # same logical [B, H, N, C] values + assert out.permute(*order).is_contiguous() # what the backend's permute(*layer_view_order) needs diff --git a/python/aisimulate/tests/unit/collector/test_vllm_collect_attn.py b/python/aisimulate/tests/unit/collector/test_vllm_collect_attn.py index 74510ddcb8..56613cf309 100644 --- a/python/aisimulate/tests/unit/collector/test_vllm_collect_attn.py +++ b/python/aisimulate/tests/unit/collector/test_vllm_collect_attn.py @@ -125,7 +125,12 @@ def __init__(self, _device): self._q_scale_float = self._k_scale_float = self._v_scale_float = 1.0 config = SimpleNamespace( - cache_config=SimpleNamespace(block_size=64, num_gpu_blocks=8192), + cache_config=SimpleNamespace( + block_size=64, + num_gpu_blocks=8192, + # resolved layout (772728f3): identity order = HND + get_resolved_kv_cache_layout=lambda: SimpleNamespace(name="LBHNC", layer_view_order=(0, 1, 2, 3)), + ), model_config=SimpleNamespace(get_sliding_window=lambda: None), ) @@ -164,6 +169,15 @@ def benchmark(**kwargs): "benchmark_with_power": benchmark, "log_perf": lambda **kwargs: calls.setdefault("log", kwargs), "vllm_version": "0.24.0", + # 1d0492eb: the page is asked from the framework per backend (collect_attn:201 + # -> utils.framework_kv_block_size); the test pins the same 16 + "framework_kv_block_size": lambda _backend_cls: 16, + # framework-resolved KV layout (collect_attn:248-270, 772728f3): identity order = HND + "get_supported_kv_cache_layouts": lambda _backends: [SimpleNamespace(name="LBHNC")], + "resolve_kv_cache_layout": lambda _cfg, _layouts, _specs: None, + "get_kv_cache_layout": lambda: SimpleNamespace(name="LBHNC", layer_view_order=(0, 1, 2, 3)), + # spy: the paged KV must be laid out in the framework-RESOLVED layout order + "kv_cache_in_resolved_layout": lambda kv, order: (calls.setdefault("kv_layout_order", order), kv)[1], } _load_function(ATTN_SOURCE, "_dense_kernel_source", namespace) run = _load_function(ATTN_SOURCE, "run_attention_torch", namespace) @@ -186,6 +200,7 @@ def benchmark(**kwargs): assert calls["log"]["item_list"][0]["isl"] == 1 assert calls["log"]["item_list"][0]["step"] == 64 assert calls["log"]["kernel_source"] == "vllm_flash_attn_fa3" + assert calls["kv_layout_order"] == (0, 1, 2, 3) # cache_config.get_resolved_kv_cache_layout() assert [event[0] for event in events] == ["update", "update", "forward", "update", "forward"] assert events[0][5] is history_slot_mapping for update, forward in zip(events[1::2], events[2::2], strict=True): diff --git a/python/aisimulate/tests/unit/collector/test_vllm_dsv4_attn.py b/python/aisimulate/tests/unit/collector/test_vllm_dsv4_attn.py index f06ae0c006..c57ddea11a 100644 --- a/python/aisimulate/tests/unit/collector/test_vllm_dsv4_attn.py +++ b/python/aisimulate/tests/unit/collector/test_vllm_dsv4_attn.py @@ -147,7 +147,7 @@ def test_smoke_has_zero_and_nonzero_prefix_without_endpoint_sweep(dsv4_module, m assert {(case[0], case[9]) for case in cases} == {(64, 0), (64, 128)} -@pytest.mark.parametrize("kernel", ["paged_mqa_logits", "hca_attn"]) +@pytest.mark.parametrize("kernel", ["paged_mqa_logits"]) def test_sparse_cases_queue_short_shapes_for_runtime_observation(dsv4_module, monkeypatch, kernel): mod = dsv4_module monkeypatch.setattr(sys, "argv", ["pytest"]) @@ -155,7 +155,6 @@ def test_sparse_cases_queue_short_shapes_for_runtime_observation(dsv4_module, mo monkeypatch.setattr(mod, "_DSV4_SPARSE_BS_LIST", [1]) monkeypatch.setattr(mod, "_DSV4_SPARSE_ISL_LIST", [1]) monkeypatch.setattr(mod, "_DSV4_SPARSE_PAST_KV_LIST", [0]) - monkeypatch.setattr(mod, "_DSV4_SPARSE_TP_LIST_ATTN", [1]) monkeypatch.setattr(mod, "_DSV4_SPARSE_TP_LIST_INDEXER", [1]) assert mod._build_vllm_dsv4_sparse_test_cases(kernel) == [[1, 1, 0, 1, kernel, "model"]] @@ -188,7 +187,7 @@ def test_dsv4_worker_rejects_ignored_runtime_inputs(dsv4_module, override, messa dsv4_module.run_dsv4_attn_worker(**args) -@pytest.mark.parametrize("kernel", ["paged_mqa_logits", "hca_attn"]) +@pytest.mark.parametrize("kernel", ["paged_mqa_logits"]) def test_sparse_worker_observes_short_shapes(dsv4_module, monkeypatch, kernel): calls = {} monkeypatch.setattr(dsv4_module, "_init_cuda", lambda device: calls.setdefault("device", device)) diff --git a/python/aisimulate/tests/unit/collector/test_vllm_kv_block_size.py b/python/aisimulate/tests/unit/collector/test_vllm_kv_block_size.py new file mode 100644 index 0000000000..f54aee205d --- /dev/null +++ b/python/aisimulate/tests/unit/collector/test_vllm_kv_block_size.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""framework_kv_block_size: the cache page comes from the model's first attention layer. + +The fakes reproduce vLLM 0.30.0's rules: AttentionBackend.get_preferred_block_size +(v1/attention/backend.py:150-158) and select_common_block_size (v1/worker/utils.py:330-385). +""" + +import ast +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + +pytestmark = pytest.mark.unit +UTILS = Path(__file__).resolve().parents[3] / "collector" / "vllm" / "utils.py" + + +def _backend(*sizes): + class Backend: + @staticmethod + def get_supported_kernel_block_sizes(): + return list(sizes) + + @classmethod + def get_preferred_block_size(cls, default_block_size): + return default_block_size if default_block_size in sizes else min(sizes) + + return Backend + + +def _select_common_block_size(kv_manager_block_size, backends): + if all(kv_manager_block_size in b.get_supported_kernel_block_sizes() for b in backends): + return kv_manager_block_size + common = set.intersection(*(set(b.get_supported_kernel_block_sizes()) for b in backends)) + return max(size for size in common if kv_manager_block_size % size == 0) + + +@pytest.fixture +def kv_block_size(monkeypatch): + for name, attrs in ( + ("vllm", {}), + ("vllm.config", {}), + ("vllm.config.cache", {"CacheConfig": SimpleNamespace(DEFAULT_BLOCK_SIZE=16)}), + ("vllm.v1", {}), + ("vllm.v1.worker", {}), + ("vllm.v1.worker.utils", {"select_common_block_size": _select_common_block_size}), + ): + module = ModuleType(name) + for key, value in attrs.items(): + setattr(module, key, value) + monkeypatch.setitem(sys.modules, name, module) + tree = ast.parse(UTILS.read_text()) + function = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "framework_kv_block_size") + namespace = {} + exec(compile(ast.Module(body=[function], type_ignores=[]), str(UTILS), "exec"), namespace) + return namespace["framework_kv_block_size"] + + +def test_dsa_sparse_mla_runs_at_the_indexer_cache_page(kv_block_size): + flashinfer_mla_sparse = _backend(32, 64) # flashinfer_mla_sparse.py:85-86 @0.30.0 + dsa_indexer = _backend(64) # mla/indexer.py:203-204 @0.30.0 (CUDA) + # alone, the sparse MLA backend would pick its own preferred 32 ... + assert kv_block_size(flashinfer_mla_sparse) == 32 + # ... but serving sizes the cache from the indexer, registered first, and keeps 64 + assert kv_block_size(flashinfer_mla_sparse, cache_backend_cls=dsa_indexer) == 64 + assert kv_block_size(flashinfer_mla_sparse, cache_backend_cls=dsa_indexer, user_block_size=32) == 32 + + +def test_first_layer_backend_is_the_default_cache_backend(kv_block_size): + flash_attn = _backend(16, 32, 64) + assert kv_block_size(flash_attn) == 16 + assert kv_block_size(_backend(64)) == 64 diff --git a/python/aisimulate/tests/unit/collector/test_vllm_mla_cli.py b/python/aisimulate/tests/unit/collector/test_vllm_mla_cli.py index f4e28658ed..925f24f041 100644 --- a/python/aisimulate/tests/unit/collector/test_vllm_mla_cli.py +++ b/python/aisimulate/tests/unit/collector/test_vllm_mla_cli.py @@ -73,6 +73,8 @@ def measure(**kwargs): names = { "_supported_model_map", "_get_precision_combos", + "_dsa_context_point_allowed", + "_dsa_context_budgets", "get_context_test_cases", "get_generation_test_cases", "run_mla_module_worker", @@ -110,10 +112,12 @@ def test_explicit_checkpoint_alias_is_preserved(cli): def test_dsa_context_cli_passes_declared_cached_prefix_to_native_worker(cli): - calls = cli("--mode", "context", "--model", "zai-org/GLM-5.2", "--gemm-type", "fp8", "--kv-cache-dtype", "fp8") + calls = cli( + "--mode", "context", "--model", "zai-org/GLM-5.2", "--gemm-type", "fp8_block", "--kv-cache-dtype", "fp8" + ) assert [(row["seq_len"], row["prefix_len"]) for row in calls] == [(128, 0), (128, 131072)] assert all(row["batch_size"] == 1 and row["num_heads"] == 8 for row in calls) - assert all(row["gemm_type"] == "fp8" and row["perf_filename"] == PerfFile.DSA_CONTEXT_MODULE for row in calls) + assert all(row["gemm_type"] == "fp8_block" and row["perf_filename"] == PerfFile.DSA_CONTEXT_MODULE for row in calls) @pytest.mark.parametrize( diff --git a/python/aisimulate/tests/unit/collector/test_vllm_mla_graph_timing.py b/python/aisimulate/tests/unit/collector/test_vllm_mla_graph_timing.py index f0880846cb..09a18de6fe 100644 --- a/python/aisimulate/tests/unit/collector/test_vllm_mla_graph_timing.py +++ b/python/aisimulate/tests/unit/collector/test_vllm_mla_graph_timing.py @@ -5,7 +5,9 @@ import ast import math +import os import sys +import time from contextlib import ExitStack, contextmanager, nullcontext from functools import wraps from pathlib import Path @@ -24,8 +26,9 @@ def _load_function(path, name, namespace): return namespace[name] -def _module_runner(*, graph_flag=True, failure=None, reuse_frequency=1, reuse_pattern=None): +def _module_runner(*, graph_flag=True, failure=None, reuse_frequency=1, reuse_pattern=None, point_failures=None): events, published, benchmark_calls = [], [], [] + point_failures = dict(point_failures or {}) error = RuntimeError(f"{failure} failed") active_contexts = [] active_slot_mapping = {} @@ -102,6 +105,18 @@ def benchmark(**kwargs): yield {"used_cuda_graph": graph_flag, "latency_ms": 0.015, "power_stats": None} events.append("graph teardown") + def create_kv(**kwargs): + events.append(("kv", kwargs["seq_len"])) + if kwargs["seq_len"] in point_failures: + raise point_failures[kwargs["seq_len"]] + return ( + Tensor(), + object(), + SimpleNamespace(slot_mapping=query_slots), + Tensor() if kwargs["attn_type"] == "dsa" else None, + object() if kwargs["attn_type"] == "dsa" else None, + ) + def cleanup(): assert active_contexts == [] assert module.mla_attn.skip_topk is False @@ -121,20 +136,17 @@ def cleanup(): set_device=lambda value: None, get_device_name=lambda value: "B200", OutOfMemoryError=MemoryError, + empty_cache=lambda: events.append("empty_cache"), ), ), + "gc": SimpleNamespace(collect=lambda: None), "setup_distributed": lambda device: None, "enable_engine_fused_ops": lambda: None, "init_workspace_manager": lambda device: events.append("workspace init"), "_create_attention_module": create_module, + "_install_rope_missing_key_shims": lambda module: ([], []), "_process_module_weights": lambda *args: None, - "_create_kv_cache_and_metadata": lambda **kwargs: ( - Tensor(), - object(), - SimpleNamespace(slot_mapping=query_slots), - Tensor() if kwargs["attn_type"] == "dsa" else None, - object() if kwargs["attn_type"] == "dsa" else None, - ), + "_create_kv_cache_and_metadata": create_kv, "set_current_vllm_config": lambda config: context("config"), "set_forward_context": forward_context, "benchmark_with_power": benchmark, @@ -142,21 +154,36 @@ def cleanup(): "log_perf": lambda **kwargs: published.append(kwargs), "_cleanup": cleanup, "vllm_version": "0.25.0", - "traceback": SimpleNamespace(print_exc=lambda: None), + "traceback": SimpleNamespace(print_exc=lambda: None, clear_frames=lambda _tb: None), + "time": time, } _load_function(COLLECTOR / "vllm/utils.py", "with_exit_stack", namespace) + # A prefix list expands through the plan predicate: budgets from the sweep spec, + # ceiling from the model config (both faked here). + namespace["get_mla_module_sweep_spec"] = lambda _backend=None: SimpleNamespace( + context_dsa_chunk_prefill_size=16384, context_max_tokens=131072, context_dsa_max_full_sequence_length=1048576 + ) + namespace["_model_max_position_embeddings"] = lambda _model: None + for helper in ( + "_series_seq_lens", + "_dsa_context_point_allowed", + "_dsa_context_budgets", + "_series_points", + "_is_fatal_cuda_error", + ): + _load_function(COLLECTOR / "vllm/collect_mla_module.py", helper, namespace) run = _load_function(COLLECTOR / "vllm/collect_mla_module.py", "run_mla_module", namespace) - def execute(attn_type="dsa", phase="context"): + def execute(attn_type="dsa", phase="context", seq_len=16, prefix_len=128): return run( - seq_len=16, + seq_len=seq_len, batch_size=2, num_heads=8, kv_cache_dtype="fp8", compute_dtype="bfloat16", gemm_type="bfloat16", perf_filename=f"{attn_type}_{phase}_module_perf.txt", - prefix_len=128, + prefix_len=prefix_len, model_path="fixture/model", attn_type=attn_type, ) @@ -169,7 +196,7 @@ def execute(attn_type="dsa", phase="context"): @pytest.mark.parametrize("reuse_frequency", [1, None]) def test_graph_publication_and_workspace_teardown_after_context_exit(attn_type, phase, reuse_frequency): execute, events, published, calls, _ = _module_runner(reuse_frequency=reuse_frequency) - assert execute(attn_type, phase) == 0.015 + assert execute(attn_type, phase) == [0.015] assert len(calls) == len(published) == 1 assert events[-4:] == ["graph teardown", "exit forward", "exit config", "cleanup"] row = published[0]["item_list"][0] @@ -185,7 +212,7 @@ def test_graph_publication_and_workspace_teardown_after_context_exit(attn_type, @pytest.mark.parametrize("reuse", [{"reuse_frequency": 4}, {"reuse_frequency": None, "reuse_pattern": "FFSS"}]) def test_reuse_checkpoint_publishes_native_full_and_skip_measurements(phase, reuse): execute, events, published, calls, _ = _module_runner(**reuse) - assert execute(phase=phase) == 0.015 + assert execute(phase=phase) == [0.015] assert [call["skip_indexer"] for call in calls] == [False, True] assert [row["op_name"] for row in published] == [ f"dsa_{phase}_module", @@ -218,6 +245,80 @@ def test_ordinary_mla_does_not_publish_dsa_reuse_rows(): assert [row["op_name"] for row in published] == ["mla_context_module"] +def test_series_builds_the_module_once_and_measures_every_point_largest_first(): + execute, events, published, calls, _ = _module_runner(reuse_frequency=4) + assert execute(seq_len=[8, 16, 8]) == [0.015, 0.015] + assert events.count("workspace init") == 1 + assert [e for e in events if isinstance(e, tuple) and e[0] == "kv"] == [("kv", 16), ("kv", 8)] + assert [(row["op_name"], row["item_list"][0]["isl"]) for row in published] == [ + ("dsa_context_module", 16), + ("dsa_context_module_skip_indexer", 16), + ("dsa_context_module", 8), + ("dsa_context_module_skip_indexer", 8), + ] + assert len(calls) == 4 and events[-1] == "cleanup" and events.count("cleanup") == 1 + + +def test_series_sweeps_the_isl_x_prefix_grid_of_one_module_largest_total_context_first(): + """A DSA context case carries its prefix list: one module, every (isl, prefix) point of the + plan grid measured on it, largest total context first (module build was 0.9-1.2 s/pt under + the campaign's 8 workers when it was amortised over one prefix only, 2026-10-29).""" + execute, events, published, calls, _ = _module_runner(reuse_frequency=4) + assert execute(seq_len=[8, 16], prefix_len=[0, 128]) == [0.015] * 4 + assert events.count("workspace init") == 1 + kv_events = [e for e in events if isinstance(e, tuple) and e[0] == "kv"] + assert kv_events == [("kv", 16), ("kv", 8), ("kv", 16), ("kv", 8)] + rows = [ + (row["item_list"][0]["isl"], row["item_list"][0]["step"]) + for row in published + if row["op_name"] == "dsa_context_module" + ] + assert rows == [(16, 128), (8, 128), (16, 0), (8, 0)] + assert len(calls) == 8 and events.count("cleanup") == 1 + + +def test_series_keeps_measuring_after_a_non_fatal_point_failure_and_raises_it_afterwards(): + oom = MemoryError("CUDA out of memory (fixture)") + execute, events, published, calls, _ = _module_runner(point_failures={16: oom}) + with pytest.raises(MemoryError) as caught: + execute(seq_len=[16, 8]) + assert caught.value is oom + assert [e for e in events if isinstance(e, tuple) and e[0] == "kv"] == [("kv", 16), ("kv", 8)] + assert "empty_cache" in events, "the failed point's memory is released before the next point" + assert [row["item_list"][0]["isl"] for row in published] == [8] + assert len(calls) == 1 and events[-1] == "cleanup" + + +def test_series_summarises_several_point_failures_in_one_classified_error(): + failures = {16: MemoryError("oom at 16"), 4: RuntimeError("Missing TRTLLM-GEN kernel (decode)")} + execute, _, published, _, _ = _module_runner(point_failures=failures) + with pytest.raises(RuntimeError, match=r"2/3 point\(s\) of the module series failed") as caught: + execute(seq_len=[16, 8, 4]) + assert "s=16 prefix=128: MemoryError: oom at 16" in str(caught.value) + assert "s=4 prefix=128: RuntimeError: Missing TRTLLM-GEN kernel" in str(caught.value) + assert [row["item_list"][0]["isl"] for row in published] == [8] + + +def test_series_aborts_at_once_on_a_sticky_cuda_error(): + sticky = RuntimeError("CUDA error: an illegal memory access was encountered") + execute, events, published, _, _ = _module_runner(point_failures={16: sticky}) + with pytest.raises(RuntimeError) as caught: + execute(seq_len=[16, 8]) + assert caught.value is sticky + kv_points = [e for e in events if isinstance(e, tuple) and e[0] == "kv"] + assert kv_points == [("kv", 16)], "no further point on a dead context" + assert published == [] and events[-1] == "cleanup" + + +def test_single_point_failures_keep_their_original_exception_and_timing(): + """A one-point case (CLI, single-point fills) behaves exactly like before the series layout.""" + oom = MemoryError("oom") + execute, events, published, _, _ = _module_runner(point_failures={16: oom}) + with pytest.raises(MemoryError) as caught: + execute(seq_len=16) + assert caught.value is oom and published == [] and events[-1] == "cleanup" + + def test_explicit_full_pattern_overrides_reuse_frequency(): execute, _, published, calls, _ = _module_runner(reuse_frequency=4, reuse_pattern="FFFF") execute() @@ -259,9 +360,11 @@ def test_indexer_output_budget_excludes_cached_prefix(monkeypatch, batch, query, """Run actual construction and inspect both tensor shape and context capacity.""" allocations, configs = [], [] - class Module: - def __init__(self, **kwargs): - assert kwargs["topk_indices_buffer"] == (expected_tokens, 2048) + class Attention: + topk_indices_buffer = "meta" # what the meta-device model allocated + + def named_modules(self): + return [("", self)] def parameters(self): return [] @@ -277,9 +380,19 @@ def eval(self): def requires_grad_(self, enabled): assert enabled is False + attention = Attention() + + class Model: + def __init__(self, *, vllm_config, prefix): + assert prefix == "" + + def named_modules(self): + return [("model.layers.0.self_attn", attention)] + def create_config(**kwargs): configs.append(kwargs) hf = SimpleNamespace( + architectures=["DeepseekV32ForCausalLM"], index_topk=2048, hidden_size=7168, qk_nope_head_dim=192, @@ -299,7 +412,12 @@ def empty(*shape, **kwargs): return shape monkeypatch.setitem( - sys.modules, "vllm.model_executor.models.deepseek_v2", SimpleNamespace(DeepseekV2MLAAttention=Module) + sys.modules, "vllm.model_executor.models.deepseek_v2", SimpleNamespace(DeepseekV2MLAAttention=object) + ) + monkeypatch.setitem( + sys.modules, + "vllm.model_executor.models.registry", + SimpleNamespace(ModelRegistry=SimpleNamespace(resolve_model_cls=lambda archs, cfg: (Model, archs[0]))), ) monkeypatch.setitem( sys.modules, @@ -311,7 +429,8 @@ def empty(*shape, **kwargs): ) namespace = { "math": math, - "torch": SimpleNamespace(empty=empty, int32="int32", no_grad=nullcontext), + "os": os, + "torch": SimpleNamespace(empty=empty, int32="int32", no_grad=nullcontext, device=lambda name: nullcontext()), "_resolve_model_path": lambda model: model, "create_vllm_config": create_config, "_create_gemm_quant_config": lambda gemm: None, @@ -331,5 +450,9 @@ def empty(*shape, **kwargs): is_context=is_context, ) assert allocations == [(expected_tokens, 2048)] + assert attention.topk_indices_buffer == (expected_tokens, 2048) # swapped in for the meta buffer assert configs[0]["max_model_len"] == max(prefix + query, 4096) - assert configs[0]["num_gpu_blocks"] == max(1 + math.ceil((prefix + query + 1) / 64) * batch, 8192) + # The module is built before the backend (and its page) is known: the KV + # capacity is an upper bound at the smallest page (16) and is recomputed + # per selected page in _create_kv_cache_and_metadata. + assert configs[0]["num_gpu_blocks"] == max(1 + math.ceil((prefix + query + 1) / 16) * batch, 8192) diff --git a/python/aisimulate/tests/unit/collector/test_vllm_mla_input_population.py b/python/aisimulate/tests/unit/collector/test_vllm_mla_input_population.py index ceb8c96599..636774d5dd 100644 --- a/python/aisimulate/tests/unit/collector/test_vllm_mla_input_population.py +++ b/python/aisimulate/tests/unit/collector/test_vllm_mla_input_population.py @@ -58,6 +58,87 @@ def test_materialization_rejects_meta_runtime_buffer_without_a_constructor_value move(module, "cpu") +class _Rotary: + """Stand-in for vLLM's RotaryEmbeddingBase: a non-persistent cos/sin table built by a + device-free recipe in __init__ (rotary_embedding/base.py:60-63 @ v0.30.0).""" + + def __new__(cls, torch, scale, dtype): + class Rotary(torch.nn.Module): + def __init__(self): + super().__init__() + self.scale = scale + self.register_buffer("cos_sin_cache", self._compute_cos_sin_cache().to(dtype), persistent=False) + + def _compute_cos_sin_cache(self): + return torch.arange(6, dtype=torch.float32).reshape(3, 2) * self.scale + + return Rotary() + + +def _meta_attention_module(torch): + # Serving-parity construction instantiates the model under torch.device("meta"): the + # two rotary tables (rotary_emb, indexer_rope_emb) and the KV-quant scale buffers that + # MLAAttention registers (attention.py:129-135,184 @ v0.30.0) are all meta. + with torch.device("meta"): + module = torch.nn.Module() + module.projection = torch.nn.Linear(4, 4) + module.rotary_emb = _Rotary(torch, 1.0, torch.bfloat16) + module.indexer_rope_emb = _Rotary(torch, 0.5, torch.float32) + module.mla_attn = torch.nn.Module() + for scale in ("_k_scale", "_v_scale", "_q_scale", "_prob_scale"): + module.mla_attn.register_buffer(scale, torch.tensor(1.0, dtype=torch.float32)) + module.mla_attn._k_scale_cpu = torch.tensor(1.0, dtype=torch.float32) + module.mla_attn._v_scale_cpu = torch.tensor(1.0, dtype=torch.float32) + return module + + +def test_materialization_rebuilds_meta_rotary_tables_and_quant_scales_from_the_owner_recipe(torch): + module = _meta_attention_module(torch) + assert module.rotary_emb.cos_sin_cache.is_meta and module.mla_attn._k_scale.is_meta + + move = _load_function("_move_module_preserving_buffers", torch) + moved = move(module, "cpu") + + recipe = torch.arange(6, dtype=torch.float32).reshape(3, 2) + assert moved is module + assert not moved.projection.weight.is_meta + assert moved.rotary_emb.cos_sin_cache.dtype == torch.bfloat16 + assert torch.equal(moved.rotary_emb.cos_sin_cache, recipe.to(torch.bfloat16)) + assert moved.indexer_rope_emb.cos_sin_cache.dtype == torch.float32 + assert torch.equal(moved.indexer_rope_emb.cos_sin_cache, recipe * 0.5) + assert "rotary_emb.cos_sin_cache" not in moved.state_dict() + for scale in ("_k_scale", "_v_scale", "_q_scale", "_prob_scale"): + tensor = getattr(moved.mla_attn, scale) + assert not tensor.is_meta and tensor.dtype == torch.float32 and tensor.item() == 1.0 + assert f"mla_attn.{scale}" in moved.state_dict() + assert moved.mla_attn._k_scale_cpu.device.type == "cpu" and moved.mla_attn._k_scale_cpu.item() == 1.0 + assert moved.mla_attn._v_scale_cpu.item() == 1.0 + + +def test_materialization_handles_a_rotary_shared_between_attention_and_indexer(torch): + # get_rope() caches instances per argument tuple, so both names can hold one module. + with torch.device("meta"): + module = torch.nn.Module() + module.projection = torch.nn.Linear(4, 4) + shared = _Rotary(torch, 2.0, torch.bfloat16) + module.rotary_emb = shared + module.indexer_rope_emb = shared + move = _load_function("_move_module_preserving_buffers", torch) + moved = move(module, "cpu") + recipe = (torch.arange(6, dtype=torch.float32).reshape(3, 2) * 2.0).to(torch.bfloat16) + assert moved.rotary_emb is moved.indexer_rope_emb + assert torch.equal(moved.indexer_rope_emb.cos_sin_cache, recipe) + + +def test_materialization_refuses_a_meta_buffer_whose_owner_has_no_recipe(torch): + module = _meta_attention_module(torch) + with torch.device("meta"): + module.register_buffer("block_table_cache", torch.zeros(4), persistent=False) + move = _load_function("_move_module_preserving_buffers", torch) + with pytest.raises(RuntimeError, match="uninitialized meta buffer.*block_table_cache"): + move(module, "cpu") + + @pytest.mark.parametrize("dtype_name", ["float8_e4m3fn", "float8_e5m2"]) def test_synthetic_fp8_weights_produce_nonzero_queries_without_changing_buffers_or_rng(dtype_name, torch): def make_module(): @@ -151,19 +232,21 @@ def test_indexer_history_rejects_an_incompatible_packed_cache(monkeypatch, torch ) -@pytest.mark.parametrize("mode,block", [("fp8", None), ("fp8_block", [128, 128])]) -def test_fp8_modes_construct_distinct_native_configs(monkeypatch, mode, block): +def test_fp8_block_constructs_the_native_block_config_and_per_tensor_fp8_is_gone(monkeypatch): monkeypatch.setitem( sys.modules, "vllm.model_executor.layers.quantization.fp8", SimpleNamespace(Fp8Config=lambda **kwargs: kwargs), ) - config = _load_function("_create_gemm_quant_config")(mode) - assert config == { + create = _load_function("_create_gemm_quant_config") + assert create("fp8_block") == { "is_checkpoint_fp8_serialized": True, "activation_scheme": "dynamic", - "weight_block_size": block, + "weight_block_size": [128, 128], } + # per-tensor fp8 MLA/DSA module combos were removed 2026-10-30 + with pytest.raises(ValueError, match="Unknown gemm_type"): + create("fp8") def test_nvfp4_config_preserves_native_modelopt_without_projection_exclusions(monkeypatch): @@ -179,9 +262,15 @@ def test_nvfp4_config_preserves_native_modelopt_without_projection_exclusions(mo } -def _nvfp4_fixture(monkeypatch, torch, *, invalid_scales=False): +def _nvfp4_fixture(monkeypatch, torch, *, invalid_scales=False, w4a16=False): + # vLLM 0.30: ModelOptLinearMethod carries a (weight, activation) QuantSpec and a + # CkptCtx; NVFP4 is recognized by spec.weight == kNvfp4Static, W4A16 has no + # activation scheme and therefore no input_scale (modelopt.py:1913-2028,2448-2476). + nvfp4_static = object() + class NativeMethod: - quant_config = SimpleNamespace(group_size=16) + spec = SimpleNamespace(weight=nvfp4_static, activation=None if w4a16 else object()) + ctx = SimpleNamespace(group_size=16) calls = [] @@ -207,16 +296,19 @@ def native_quantize(values, inverse_scale, *, is_sf_swizzled_layout): ) monkeypatch.setitem( sys.modules, - "vllm.model_executor.layers.quantization.modelopt", - SimpleNamespace(ModelOptNvFp4LinearMethod=NativeMethod), + "vllm.model_executor.layers.quantization.utils.quant_utils", + SimpleNamespace(kNvfp4Static=nvfp4_static), ) module = torch.nn.Module() layer = module.projection = torch.nn.Module() layer.quant_method = NativeMethod() layer.weight = torch.nn.Parameter(torch.zeros(1025, 512, dtype=torch.uint8), requires_grad=False) - layer.weight_scale = torch.nn.Parameter(torch.zeros(1025, 64, dtype=torch.float8_e4m3fn), requires_grad=False) + layer.weight_scale = torch.nn.Parameter( + torch.full((1025, 64), float("nan")).to(torch.float8_e4m3fn), requires_grad=False + ) layer.weight_scale_2 = torch.nn.Parameter(torch.zeros(2), requires_grad=False) - layer.input_scale = torch.nn.Parameter(torch.zeros(2), requires_grad=False) + if not w4a16: + layer.input_scale = torch.nn.Parameter(torch.zeros(2), requires_grad=False) module.other_weight = torch.nn.Parameter(torch.full((2, 2), 17.0), requires_grad=False) module.register_buffer("rotary", torch.tensor([0.3, 0.7]), persistent=False) return module, calls @@ -257,3 +349,14 @@ def test_nvfp4_cannot_be_claimed_without_a_native_nvfp4_projection(monkeypatch, module.projection.quant_method = object() with pytest.raises(RuntimeError, match="no native ModelOpt NVFP4 linear"): _load_function("_initialize_nvfp4_parameters", torch)(module) + + +def test_nvfp4_w4a16_projection_is_populated_without_an_input_scale(monkeypatch, torch): + module, calls = _nvfp4_fixture(monkeypatch, torch, w4a16=True) + _load_function("_initialize_nvfp4_parameters", torch)(module) + layer = module.projection + assert not hasattr(layer, "input_scale") + assert [v.shape[0] for v, _ in calls] == [1024, 1] + assert torch.all(layer.weight == 0x12) + assert not torch.isnan(layer.weight_scale.float()).any() + assert torch.unique(layer.weight_scale_2).numel() == 1 and torch.all(layer.weight_scale_2 > 0) diff --git a/python/aisimulate/tests/unit/collector/trtllm/test_collect_dsv4_attn.py b/python/aisimulate/tests/unit/collector/trtllm/test_collect_dsv4_attn.py index 7c5cf83e6f..9b8fc6fc1f 100644 --- a/python/aisimulate/tests/unit/collector/trtllm/test_collect_dsv4_attn.py +++ b/python/aisimulate/tests/unit/collector/trtllm/test_collect_dsv4_attn.py @@ -36,6 +36,8 @@ def _load_module_with_torch_stub(monkeypatch): assert spec and spec.loader module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) + # no GPU under unit tests: the population / serving-KV-dtype paths read the SM + monkeypatch.setattr(module, "get_sm_version", lambda: 90, raising=False) return module @@ -172,7 +174,7 @@ def test_module_cache_reuses_same_geometry_and_evicts_on_change(monkeypatch): builds = [] - def fake_build(*, model_path, attn_kind, tp_size, device): + def fake_build(*, model_path, attn_kind, tp_size, device, kv_cache_dtype="fp8"): builds.append((model_path, attn_kind, tp_size)) return (object(), object(), {"local_heads": 8, "native_heads": 64}) @@ -186,7 +188,7 @@ def fake_build(*, model_path, attn_kind, tp_size, device): b1 = module._cached_dsv4_attention_module("m/flash", "hca", 4, "cuda:0") assert len(builds) == 2 - assert list(module._MODULE_CACHE) == [("m/flash", "hca", 4, "cuda:0")] + assert list(module._MODULE_CACHE) == [("m/flash", "hca", 4, "cuda:0", "fp8")] # key carries the KV dtype (189589d2) assert b1 is not a1 module._MODULE_CACHE.clear() @@ -271,6 +273,8 @@ def _kv_cache_config(**kw): return SimpleNamespace(**kw) llm_args.KvCacheConfig = _kv_cache_config + # no GPU under unit tests: pin the Hopper branch (fp8_ds_mla pool) deterministically + monkeypatch.setattr(module, "get_sm_version", lambda: 90) # the create function re-imports this locally (relative-or-plain fallback) mla_stub = _mod("collect_mla_module") mla_stub.get_kv_cache_manager_cls = lambda mc, cfg: _CapturingManager @@ -460,3 +464,48 @@ def test_constructor_args_match_the_serving_oracle(monkeypatch): # position == past_seen (the runner builds the tensor from the # single-sourced geometry helper) assert module.generation_request_geometry(sl)["position"] == oracle["first_position"] + + +@pytest.mark.parametrize( + "bs, sl, is_ctx, prefix", + [(1, 1, True, 0), (16, 10240, True, 4096), (128, 65536, False, 0)], +) +def test_kv_quota_is_capped_by_this_batch_s_tokens(monkeypatch, bs, sl, is_ctx, prefix): + """Per-case KV pool: the byte quota stays the serving half-of-free bound, + AND a max_tokens cap sized to this batch (doubled, page-rounded) is handed + to the manager's own DSV4 arithmetic — never collector-computed bytes.""" + module = _load_module_with_torch_stub(monkeypatch) + _, _, captured = _run_create(module, monkeypatch, batch_size=bs, seq_len=sl, is_context=is_ctx, prefix_len=prefix) + (kw,) = captured["kv_cache_config"] + assert kw["max_gpu_total_bytes"] == 50 * 2**30 # mem_get_info stub: free=100 GiB + tpb = kw["tokens_per_block"] + envelope = (prefix + sl + 1) if is_ctx else (sl + 1) + envelope = max(envelope, 512) + needed = bs * envelope + assert kw["max_tokens"] % tpb == 0 + floor = module._KV_POOL_MIN_TOKENS + expected = max(2 * needed, floor) + assert expected <= kw["max_tokens"] <= max(2 * (needed + bs * tpb) + tpb, floor) + # the floor keeps the footer-scale pool past inductor's int64 threshold + assert floor > 1_572_864 # H20 calibration: 1,048,576 tokens -> i32, 1,572,864 -> i64 + + +@pytest.mark.parametrize(("sm", "expected"), [(100, "bfloat16"), (90, "fp8"), (120, "fp8")]) +def test_serving_kv_cache_dtype_is_an_sdk_enum_member_name(monkeypatch, sm, expected): + # engine_table_view rehydrates kv_cache_dtype through KVCacheQuantMode; "bf16" raised + # KeyError on every SM100 DSV4 table in Full CI (2026-10-29). + from aisimulate_core.quantization import KVCacheQuantMode + + module = _load_module_with_torch_stub(monkeypatch) + monkeypatch.delenv("AIS_DSV4_KV_CACHE_DTYPE", raising=False) + monkeypatch.setattr(module, "get_sm_version", lambda: sm) + label = module._serving_kv_cache_dtype() + assert label == expected + assert label in KVCacheQuantMode.__members__ + + +@pytest.mark.parametrize(("override", "expected"), [("bf16", "bfloat16"), ("bfloat16", "bfloat16"), ("fp8", "fp8")]) +def test_serving_kv_cache_dtype_override_normalises_the_bf16_alias(monkeypatch, override, expected): + module = _load_module_with_torch_stub(monkeypatch) + monkeypatch.setenv("AIS_DSV4_KV_CACHE_DTYPE", override) + assert module._serving_kv_cache_dtype() == expected diff --git a/python/aisimulate/tests/unit/collector/trtllm/test_collect_gdn_fail_closed.py b/python/aisimulate/tests/unit/collector/trtllm/test_collect_gdn_fail_closed.py new file mode 100644 index 0000000000..406abb3d1b --- /dev/null +++ b/python/aisimulate/tests/unit/collector/trtllm/test_collect_gdn_fail_closed.py @@ -0,0 +1,147 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The trtllm GDN collector publishes CUDA-graph-timed rows only. + +Until 2026-10-11 every benchmark_with_power call in collector/trtllm/collect_gdn.py +passed allow_graph_fail=True: a failed capture logged one WARNING, timed the +Python launch loop eagerly and shipped that floor as the kernel latency with no +trace in the row. Capture failures now raise, the failed cell is counted, and +the case ends as a classified strict-completeness failure (the sglang GDN +collector's contract).""" + +import ast +from contextlib import contextmanager +from pathlib import Path +from types import SimpleNamespace + +import pytest + +pytestmark = pytest.mark.unit + +REPO_ROOT = Path(__file__).resolve().parents[4] +SOURCE = REPO_ROOT / "collector" / "trtllm" / "collect_gdn.py" + + +def _tree(): + return ast.parse(SOURCE.read_text(encoding="utf-8"), filename=str(SOURCE)) + + +def _function(name): + return next(node for node in _tree().body if isinstance(node, ast.FunctionDef) and node.name == name) + + +def test_trtllm_gdn_benchmarks_never_fall_back_to_eager_timing(): + calls = [ + node + for node in ast.walk(_tree()) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "benchmark_with_power" + ] + assert calls + # allow_graph_fail is covered repo-wide by test_no_eager_graph_fallback_sites.py + for call in calls: + assert all(kw.arg != "use_cuda_graph" for kw in call.keywords) + + +# The context runner's final raise is exercised behaviourally below. +@pytest.mark.parametrize("function_name", ["run_gdn_generation_benchmark"]) +def test_trtllm_gdn_phase_runners_end_in_a_strict_completeness_raise(function_name): + function = _function(function_name) + # A per-cell `continue` is only acceptable because the function re-raises at the end. + raises = [ + node + for node in function.body + if isinstance(node, ast.If) and any(isinstance(statement, ast.Raise) for statement in node.body) + ] + assert raises, f"{function_name} swallows cell failures without a final raise" + assert "strict completeness" in ast.unparse(raises[-1]) + + +class _Tensor: + def __init__(self, shape): + self.shape = tuple(shape) + + def reshape(self, *shape): + return _Tensor(shape) + + def float(self): + return self + + +def _fake_torch(): + def factory(*shape, **_kwargs): + return _Tensor(shape) + + return SimpleNamespace( + bfloat16="bf16", + float32="fp32", + int64="i64", + device=lambda name: name, + cuda=SimpleNamespace( + set_device=lambda _d: None, + synchronize=lambda: None, + empty_cache=lambda: None, + get_device_name=lambda _d: "fake-gpu", + ), + set_default_device=lambda _d: None, + randn=factory, + zeros=factory, + arange=lambda *a, **k: _Tensor((1,)), + sigmoid=lambda t: t, + nn=SimpleNamespace(functional=SimpleNamespace(logsigmoid=lambda t: t)), + ) + + +def _load_context_runner(benchmark_with_power, log_perf): + namespace = { + "torch": _fake_torch(), + "gc": SimpleNamespace(collect=lambda: None), + "aic_debug": 0, + "aic_cached_inputs": 1, + "benchmark_with_power": benchmark_with_power, + "log_perf": log_perf, + "causal_conv1d_fn": lambda *a, **k: None, + "chunk_gated_delta_rule": lambda *a, **k: None, + "_CHUNK_KERNEL_SOURCE": "chunk_gated_delta_rule_flashinfer", + "_make_input_pool": None, + } + module = ast.Module(body=[_function("run_gdn_context_benchmark")], type_ignores=[]) + exec(compile(module, str(SOURCE), "exec"), namespace) + return namespace["run_gdn_context_benchmark"] + + +def test_trtllm_gdn_context_capture_failure_fails_the_case_but_keeps_the_other_cells(): + calls = {"bench": 0} + logged = [] + + @contextmanager + def benchmark_with_power(**_kwargs): + calls["bench"] += 1 + if calls["bench"] == 3: # second (batch, seq) cell: the conv capture fails + raise RuntimeError("CUDA graph capture failed: operation not permitted when stream is capturing") + yield {"latency_ms": 0.01, "power_stats": None, "used_cuda_graph": True} + + def log_perf(**kwargs): + logged.append(kwargs["kernel_source"]) + return True + + run = _load_context_runner(benchmark_with_power, log_perf) + expected = r"strict completeness: ok=1 error=1 skip=0.*seq_len=64: RuntimeError" + with pytest.raises(RuntimeError, match=expected) as info: + run( + d_model=2048, + d_conv=4, + num_k_heads=16, + head_k_dim=128, + num_v_heads=32, + head_v_dim=128, + batch_size_list=[1], + seq_len_list=[32, 64], + model_name="Qwen/Qwen3.5-test", + perf_filename="gdn_perf.txt", + trtllm_version="1.3.0rc29", + device="cuda:0", + ) + assert "CUDA graph capture failed" in str(info.value) + # the first cell shipped both of its rows; the failed cell shipped none + assert logged == ["causal_conv1d_fn", "chunk_gated_delta_rule_flashinfer"] diff --git a/python/aisimulate/tests/unit/collector/trtllm/test_collect_gemm_op_cache.py b/python/aisimulate/tests/unit/collector/trtllm/test_collect_gemm_op_cache.py new file mode 100644 index 0000000000..81345e0749 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/trtllm/test_collect_gemm_op_cache.py @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""GPU-free tests for the TRT-LLM GEMM collector's per-(gemm_type, n, k) module cache. + +``collector/trtllm/collect_gemm.py`` imports torch and tensorrt_llm at module +scope, so ``_gemm_ops`` / ``run_gemm`` are AST-extracted and executed against +fakes (same technique as ``test_collect_moe_ep.py``). +""" + +from __future__ import annotations + +import ast +import math +from contextlib import contextmanager +from pathlib import Path +from types import SimpleNamespace +from typing import ClassVar + +import pytest + +pytestmark = pytest.mark.unit + +REPO_ROOT = Path(__file__).resolve().parents[4] +SOURCE_PATH = REPO_ROOT / "collector" / "trtllm" / "collect_gemm.py" +NAMES = ("_gemm_op_cache", "_gemm_ops", "run_gemm") +L2_BYTES = 126 * 2**20 + + +class _Tensor: + def __init__(self, shape): + self.shape = shape + + def to(self, *args, **kwargs): + return self + + +class _Linear: + instances: ClassVar[list] = [] + + def __init__(self, in_features, out_features, **kwargs): + self.in_features = in_features + self.out_features = out_features + self.kwargs = kwargs + self.loaded = [] + self.post_loaded = 0 + self.forwards: list = [] + _Linear.instances.append(self) + + def load_weights(self, weights): + self.loaded.append(weights) + + def post_load_weights(self): + self.post_loaded += 1 + + def to(self, device): + self.device = device + return self + + def forward(self, x): + self.forwards.append(x) + + +def _load(): + tree = ast.parse(SOURCE_PATH.read_text(), filename=str(SOURCE_PATH)) + + def wanted(node) -> bool: + if isinstance(node, ast.FunctionDef): + return node.name in NAMES + if isinstance(node, ast.AnnAssign): + return isinstance(node.target, ast.Name) and node.target.id in NAMES + return False + + selected = [node for node in tree.body if wanted(node)] + assert len(selected) == len(NAMES) + _Linear.instances = [] + built: list = [] + empty_cache_calls: list = [] + benchmarks: list = [] + published: list = [] + + @contextmanager + def benchmark_with_power(*, device, kernel_func, repeat_n): + benchmarks.append({"forwards_before": [len(op.forwards) for op in _Linear.instances], "repeat_n": repeat_n}) + kernel_func() + yield {"latency_ms": 1.0, "power_stats": None} + + namespace = { + "math": math, + "torch": SimpleNamespace( + device=lambda value: value, + bfloat16="bf16", + randn=lambda shape, dtype: _Tensor(shape), + set_default_device=lambda device: None, + cuda=SimpleNamespace( + set_device=lambda device: None, + empty_cache=lambda: empty_cache_calls.append(True), + get_device_name=lambda device: "B200", + ), + ), + "Linear": _Linear, + "_build_weights": lambda gemm_type, n, k, device, dtype, group_size, x: ( + built.append((gemm_type, n, k, group_size)) or {"weight": (n, k)} + ), + "_get_l2_cache_bytes": lambda index: L2_BYTES, + "get_sm_version": lambda: 100, + "is_sm_100f": lambda sm: sm in (100, 103), + "QuantConfig": lambda **kwargs: kwargs, + "QuantAlgo": SimpleNamespace(FP8="fp8", FP8_BLOCK_SCALES="fp8_block", NVFP4="nvfp4"), + "benchmark_with_power": benchmark_with_power, + "log_perf": lambda **kwargs: published.append(kwargs), + "tensorrt_llm": SimpleNamespace(__version__="1.3.0rc29"), + } + exec(compile(ast.Module(body=selected, type_ignores=[]), str(SOURCE_PATH), "exec"), namespace) + return namespace, built, empty_cache_calls, benchmarks, published + + +def _ops(namespace, gemm_type, n, k, *, copies): + return namespace["_gemm_ops"]( + gemm_type, + n, + k, + quant_config=None, + group_size=None, + outside_loop_count=copies, + device=SimpleNamespace(index=0), + dtype="bf16", + x=_Tensor((1, k)), + ) + + +def test_same_key_reuses_the_built_modules_and_a_new_key_replaces_them(): + namespace, built, empty_cache_calls, _, _ = _load() + first = _ops(namespace, "nvfp4", 4096, 4096, copies=5) + assert len(first) == 5 and len(built) == 5 and len(_Linear.instances) == 5 + assert empty_cache_calls == [True], "the previous key's modules are dropped before building" + assert _ops(namespace, "nvfp4", 4096, 4096, copies=5) is first + assert len(built) == 5, "a second m for the same (gemm_type, n, k) builds nothing" + second = _ops(namespace, "nvfp4", 8192, 4096, copies=2) + assert second is not first and len(second) == 2 + assert len(built) == 7 and len(empty_cache_calls) == 2 + assert list(namespace["_gemm_op_cache"]) == [("nvfp4", 8192, 4096, 2)], "one key is held at a time" + + +def test_run_gemm_builds_once_per_key_dry_runs_every_copy_and_divides_the_latency(): + namespace, built, _, benchmarks, published = _load() + n = k = 1024 # 2 MiB bf16 weight -> 5 copies to span the 126 MiB L2 + for m in (8192, 4096, 1): + namespace["run_gemm"]("bfloat16", m, n, k, perf_filename="gemm_perf.txt", device="cuda:0") + assert len(built) == 5 and len(_Linear.instances) == 5 + expected_copies = min(5, math.ceil(L2_BYTES / (n * k * 2))) + assert expected_copies == 5 + # Before each capture every cached copy saw its dry run at the current m (JIT / lazy init). + assert [b["forwards_before"] for b in benchmarks] == [[1] * 5, [3] * 5, [5] * 5] + assert all(b["repeat_n"] == 1 for b in benchmarks) + assert [row["item_list"][0]["m"] for row in published] == [8192, 4096, 1] + assert all(row["item_list"][0]["latency"] == pytest.approx(1.0 / 5) for row in published) + assert all(row["kernel_source"] == "torch_flow" for row in published) + + +def test_fp8_block_post_load_runs_once_per_copy_and_is_tagged_deepgemm_on_sm100(): + namespace, built, _, _, published = _load() + n, k = 16384, 7168 # 112 MiB fp8 weight -> 2 copies + namespace["run_gemm"]("fp8_block", 64, n, k, perf_filename="gemm_perf.txt", device="cuda:0") + namespace["run_gemm"]("fp8_block", 32, n, k, perf_filename="gemm_perf.txt", device="cuda:0") + assert built == [("fp8_block", n, k, 128)] * 2 + assert [op.post_loaded for op in _Linear.instances] == [1, 1] + expected_quant = {"quant_algo": "fp8_block", "group_size": 128} + assert [op.kwargs["quant_config"] for op in _Linear.instances] == [expected_quant] * 2 + assert [row["kernel_source"] for row in published] == ["deepgemm", "deepgemm"] + assert all(row["item_list"][0]["latency"] == pytest.approx(0.5) for row in published) diff --git a/python/aisimulate/tests/unit/collector/trtllm/test_collect_mamba2_fail_closed.py b/python/aisimulate/tests/unit/collector/trtllm/test_collect_mamba2_fail_closed.py new file mode 100644 index 0000000000..b433c1a503 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/trtllm/test_collect_mamba2_fail_closed.py @@ -0,0 +1,308 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The trtllm Mamba2 collector publishes CUDA-graph-timed rows only, never +drops a kernel lane silently, and runs the SSM kernels the way serving does. + +Until 2026-10-12 every benchmark_with_power call in collector/trtllm/collect_mamba2.py +passed allow_graph_fail=True, each (batch, seq) cell sat in a bare +`except Exception: continue`, and neither SSM kernel received the caller-owned +`out` buffer TRT-LLM requires (ssd_chunk_scan.py:667, selective_state_update.py +`out.dim()` @1.3.0rc20/rc29). The shipped rc20 mamba2 tables of all eight +systems and the rc29 b200 package contain only the causal_conv1d lanes: the +mamba_chunk_scan_combined / selective_state_update lanes vanished with zero +failure records. Capture failures now raise, every failed cell is counted, the +case ends as a classified strict-completeness failure (the GDN contract), both +SSM calls carry `out=`, and the decode state update is the kernel serving's +Mamba2Mixer gate selects (FlashInfer or TRT-LLM's Triton), labelled as such.""" + +import ast +from contextlib import contextmanager +from pathlib import Path +from types import SimpleNamespace + +import pytest + +pytestmark = pytest.mark.unit + +REPO_ROOT = Path(__file__).resolve().parents[4] +SOURCE = REPO_ROOT / "collector" / "trtllm" / "collect_mamba2.py" + + +def _tree(): + return ast.parse(SOURCE.read_text(encoding="utf-8"), filename=str(SOURCE)) + + +def _function(name): + return next(node for node in _tree().body if isinstance(node, ast.FunctionDef) and node.name == name) + + +def _callee(node): + return ast.unparse(node.func) + + +def test_trtllm_mamba2_benchmarks_never_fall_back_to_eager_timing(): + calls = [ + node for node in ast.walk(_tree()) if isinstance(node, ast.Call) and _callee(node) == "benchmark_with_power" + ] + assert calls + # allow_graph_fail is covered repo-wide by test_no_eager_graph_fallback_sites.py + for call in calls: + assert all(keyword.arg != "use_cuda_graph" for keyword in call.keywords) + + +def test_every_ssm_kernel_call_hands_over_the_caller_owned_out_buffer(): + # TRT-LLM's chunk scan asserts on `out` and its state update calls + # `out.dim()`; the missing keyword is what emptied the SSM lanes. + scan_calls = [ + node + for node in ast.walk(_function("run_mamba2_context_benchmark")) + if isinstance(node, ast.Call) and _callee(node) == "mamba_chunk_scan_combined" + ] + update_calls = [ + node + for node in ast.walk(_function("run_mamba2_generation_benchmark")) + if isinstance(node, ast.Call) and _callee(node) == "ssu_func" + ] + assert len(scan_calls) == 4 # cached + pooled, warm-up call + timed closure + assert len(update_calls) == 4 + for call in scan_calls + update_calls: + assert any(keyword.arg == "out" for keyword in call.keywords), ast.unparse(call) + for call in update_calls: + assert any(keyword.arg == "state_batch_indices" for keyword in call.keywords), ast.unparse(call) + + +def _load_gate(): + tree = _tree() + body = [ + node + for node in tree.body + if (isinstance(node, ast.Assign) and ast.unparse(node.targets[0]).startswith("_FLASHINFER_SSU_")) + or (isinstance(node, ast.FunctionDef) and node.name == "_serving_selects_flashinfer_ssu") + ] + namespace = {} + exec(compile(ast.Module(body=body, type_ignores=[]), str(SOURCE), "exec"), namespace) + return namespace["_serving_selects_flashinfer_ssu"] + + +@pytest.mark.parametrize( + ("nheads", "n_groups", "head_dim", "d_state", "flashinfer"), + [ + (64, 8, 64, 128, True), # Nemotron-3 Nano / generic 4K: ratio 8 + (256, 8, 64, 256, True), # Nemotron-H 56B: ratio 32, d_state 256 + (64, 8, 80, 128, False), # Nemotron-v2-Nano head_dim 80: no FlashInfer kernel + (48, 16, 64, 128, False), # ratio 3 + (64, 8, 64, 32, False), # d_state 32 + (64, 0, 64, 128, False), # no groups + ], +) +def test_decode_kernel_gate_mirrors_the_rc29_mixer(nheads, n_groups, head_dim, d_state, flashinfer): + assert _load_gate()(nheads, n_groups, head_dim, d_state) is flashinfer + + +class _Tensor: + """Shape-only stand-in; arithmetic and method calls return another stand-in.""" + + def __init__(self, shape=()): + self.shape = tuple(shape) + + def __getattr__(self, name): + if name.startswith("__"): + raise AttributeError(name) + return lambda *args, **kwargs: _Tensor(self.shape) + + def __neg__(self): + return _Tensor(self.shape) + + def __add__(self, other): + return _Tensor(self.shape) + + __radd__ = __add__ + __sub__ = __add__ + __rsub__ = __add__ + __mul__ = __add__ + __rmul__ = __add__ + + +def _factory(*shape, **_kwargs): + return _Tensor(shape) + + +def _fake_torch(): + return SimpleNamespace( + bfloat16="bf16", + float32="fp32", + int32="i32", + device=lambda name: name, + set_default_device=lambda _device: None, + cuda=SimpleNamespace( + set_device=lambda _device: None, + synchronize=lambda: None, + empty_cache=lambda: None, + get_device_name=lambda _device: "fake-gpu", + ), + randn=_factory, + rand=_factory, + zeros=_factory, + arange=lambda n, **_kwargs: _Tensor((n,)), + empty_like=lambda tensor: _Tensor(tensor.shape), + ) + + +def _fake_input_pool(shapes, count, _dtype, _device): + return {name: [_Tensor(shape) for _ in range(count)] for name, shape in shapes.items()} + + +def _load_runner( + name, benchmark_with_power, log_perf, *, cached_inputs=1, ssu_label="selective_state_update_flashinfer" +): + ssu_calls = [] + + def ssu(*args, **kwargs): + ssu_calls.append(kwargs) + + namespace = { + "torch": _fake_torch(), + "gc": SimpleNamespace(collect=lambda: None), + "aic_debug": 0, + "aic_cached_inputs": cached_inputs, + "benchmark_with_power": benchmark_with_power, + "log_perf": log_perf, + "repeat": lambda tensor, _pattern, **_axes: tensor, + "tensorrt_llm": SimpleNamespace(__version__="1.3.0rc29"), + "causal_conv1d_fn": lambda *args, **kwargs: None, + "causal_conv1d_update": lambda *args, **kwargs: None, + "mamba_chunk_scan_combined": lambda *args, **kwargs: None, + "selective_state_update": lambda *args, **kwargs: None, + "_decode_state_update_kernel": lambda *geometry: (ssu, ssu_label), + "_decode_dt": lambda batch_size, nheads, head_dim, _device: _Tensor((batch_size, nheads, head_dim)), + "_token_major_prefill_conv": lambda: False, + "_prefill_conv_token_major_inputs": None, + "_run_prefill_conv_token_major": None, + "_make_input_pool": _fake_input_pool, + } + module = ast.Module(body=[_function(name)], type_ignores=[]) + exec(compile(module, str(SOURCE), "exec"), namespace) + return namespace[name], ssu_calls + + +def _failing_on(call_index): + """A fake benchmark_with_power whose ``call_index``-th capture fails (0 = never).""" + calls = {"bench": 0} + logged = [] + + @contextmanager + def benchmark_with_power(**kwargs): + calls["bench"] += 1 + if calls["bench"] == call_index: + raise RuntimeError("CUDA graph capture failed: operation not permitted when stream is capturing") + kwargs["kernel_func"]() # the real helper runs the closure for its warmups and the capture + yield {"latency_ms": 0.01, "power_stats": None, "used_cuda_graph": True} + + def log_perf(**kwargs): + logged.append(kwargs["kernel_source"]) + return True + + return benchmark_with_power, log_perf, logged + + +_LAYER = { + "d_model": 4096, + "d_state": 128, + "d_conv": 4, + "nheads": 128, + "head_dim": 64, + "n_groups": 8, + "chunk_size": 256, + "model_name": "nvidia/Nemotron-test", + "perf_filename": "mamba2_perf.txt", + "device": "cuda:0", +} + + +def test_trtllm_mamba2_context_scan_capture_failure_fails_the_case_instead_of_dropping_the_lane(): + # The shipped-data failure mode: a cell's conv row ships, then its scan lane raises. + benchmark_with_power, log_perf, logged = _failing_on(2) + run, _ = _load_runner("run_mamba2_context_benchmark", benchmark_with_power, log_perf) + expected = r"strict completeness: ok=1 error=1 skip=0.*seq_len=32: RuntimeError" + with pytest.raises(RuntimeError, match=expected) as info: + run(batch_size_list=[1], seq_len_list=[32, 64], **_LAYER) + assert "CUDA graph capture failed" in str(info.value) + # cell (1, 32) shipped its conv row and lost the scan; cell (1, 64) shipped both + assert logged == ["causal_conv1d_fn", "causal_conv1d_fn", "mamba_chunk_scan_combined"] + + +@pytest.mark.parametrize("cached_inputs", [1, 0]) +def test_trtllm_mamba2_context_all_cells_captured_returns_quietly(cached_inputs): + benchmark_with_power, log_perf, logged = _failing_on(0) + run, _ = _load_runner("run_mamba2_context_benchmark", benchmark_with_power, log_perf, cached_inputs=cached_inputs) + assert run(batch_size_list=[1, 2], seq_len_list=[32], **_LAYER) is None + assert logged == ["causal_conv1d_fn", "mamba_chunk_scan_combined"] * 2 + + +def test_trtllm_mamba2_generation_state_update_capture_failure_fails_the_case(): + benchmark_with_power, log_perf, logged = _failing_on(4) + run, _ = _load_runner("run_mamba2_generation_benchmark", benchmark_with_power, log_perf) + expected = r"strict completeness: ok=1 error=1 skip=0.*batch_size=2: RuntimeError" + with pytest.raises(RuntimeError, match=expected): + run(batch_size_list=[1, 2], **_LAYER) + assert logged == ["causal_conv1d_update", "selective_state_update_flashinfer", "causal_conv1d_update"] + + +@pytest.mark.parametrize("cached_inputs", [1, 0]) +@pytest.mark.parametrize("ssu_label", ["selective_state_update_flashinfer", "selective_state_update_triton"]) +def test_trtllm_mamba2_generation_labels_the_rows_with_the_kernel_serving_selected(cached_inputs, ssu_label): + benchmark_with_power, log_perf, logged = _failing_on(0) + run, ssu_calls = _load_runner( + "run_mamba2_generation_benchmark", + benchmark_with_power, + log_perf, + cached_inputs=cached_inputs, + ssu_label=ssu_label, + ) + assert run(batch_size_list=[1], **_LAYER) is None + assert logged == ["causal_conv1d_update", ssu_label] + # warm-up call + the timed closure; every call carries the serving-shaped extras + assert len(ssu_calls) == 2 + for kwargs in ssu_calls: + assert kwargs["out"].shape == (1, 128, 64) + assert kwargs["state_batch_indices"].shape == (1,) + assert kwargs["dt_softplus"] is True and kwargs["z"] is None + + +def test_generation_runner_builds_dt_only_through_the_serving_layout_helper(): + # Both input modes take dt from _decode_dt; no dense (batch, nheads, head_dim) dt is made. + runner = _function("run_mamba2_generation_benchmark") + helper_calls = [node for node in ast.walk(runner) if isinstance(node, ast.Call) and _callee(node) == "_decode_dt"] + assert len(helper_calls) == 2 + dt_assignments = [ + node + for node in ast.walk(runner) + if isinstance(node, ast.Assign) and any(isinstance(t, ast.Name) and t.id == "dt" for t in node.targets) + ] + assert dt_assignments, "cached-input mode no longer binds dt" + for assignment in dt_assignments: + assert isinstance(assignment.value, ast.Call) and _callee(assignment.value) == "_decode_dt" + + +def test_decode_dt_is_an_fp32_per_head_tensor_broadcast_over_head_dim(): + # FlashInfer's binding requires dt.stride(1) == 1 and dt.stride(2) == 0 + # (csrc/selective_state_update.cu:187-188 @v0.6.18); serving gets there with + # fp32 (batch, nheads) + einops.repeat "b h -> b h p" (mamba2_mixer.py:764 @rc29). + seen = {} + + def randn(*shape, dtype=None, device=None): + seen["randn"] = (shape, dtype) + return _Tensor(shape) + + def repeat(tensor, pattern, **axes): + seen["repeat"] = (tensor.shape, pattern, axes) + return _Tensor(tensor.shape + (axes["p"],)) + + namespace = {"torch": SimpleNamespace(randn=randn, float32="fp32", Tensor=object), "repeat": repeat} + module = ast.Module(body=[_function("_decode_dt")], type_ignores=[]) + exec(compile(module, str(SOURCE), "exec"), namespace) + dt = namespace["_decode_dt"](4, 128, 64, "cuda:0") + assert seen["randn"] == ((4, 128), "fp32") + assert seen["repeat"] == ((4, 128), "b h -> b h p", {"p": 64}) + assert dt.shape == (4, 128, 64) diff --git a/python/aisimulate/tests/unit/collector/trtllm/test_collect_module_layer_cache.py b/python/aisimulate/tests/unit/collector/trtllm/test_collect_module_layer_cache.py new file mode 100644 index 0000000000..3a9a0e7e47 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/trtllm/test_collect_module_layer_cache.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""GPU-free tests for the per-worker attention-layer caches of the TRT-LLM MLA/DSA and +MSA module collectors (same pattern as collect_dsv4_attn._cached_dsv4_attention_module). +The modules import tensorrt_llm at module scope, so the helpers are AST-extracted.""" + +from __future__ import annotations + +import ast +from pathlib import Path +from types import SimpleNamespace + +import pytest + +pytestmark = pytest.mark.unit +COLLECTOR = Path(__file__).resolve().parents[4] / "collector" / "trtllm" + + +def _load(filename: str, helper: str, builder: str): + tree = ast.parse((COLLECTOR / filename).read_text()) + selected = [ + n + for n in tree.body + if (isinstance(n, ast.FunctionDef) and n.name == helper) + or (isinstance(n, ast.AnnAssign) and isinstance(n.target, ast.Name) and n.target.id == "_LAYER_CACHE") + ] + assert len(selected) == 2 + builds, evictions = [], [] + namespace = { + builder: lambda **kwargs: (builds.append(kwargs) or ("module", "config", "arch", len(builds))), + "gc": SimpleNamespace(collect=lambda: evictions.append("gc")), + "torch": SimpleNamespace(cuda=SimpleNamespace(empty_cache=lambda: evictions.append("empty_cache"))), + } + exec(compile(ast.Module(body=selected, type_ignores=[]), filename, "exec"), namespace) + return namespace[helper], builds, evictions + + +def test_mla_layer_is_built_once_per_key_and_evicted_on_a_new_key(): + cached, builds, evictions = _load("collect_mla_module.py", "_cached_attention_layer", "create_attention_layer") + first = cached("deepseek-ai/DeepSeek-V3.2", 128, True, "fp8_block", "cuda:0") + assert cached("deepseek-ai/DeepSeek-V3.2", 128, True, "fp8_block", "cuda:0") is first + assert len(builds) == 1 and evictions == [] + assert builds[0] == { + "model_path": "deepseek-ai/DeepSeek-V3.2", + "num_heads": 128, + "use_fp8_kv_cache": True, + "gemm_type": "fp8_block", + "device": "cuda:0", + } + second = cached("deepseek-ai/DeepSeek-V3.2", 16, True, "fp8_block", "cuda:0") + assert second is not first and len(builds) == 2 + assert evictions == ["gc", "empty_cache"], "the previous layer is released before the next build" + third = cached("deepseek-ai/DeepSeek-V3.2", 128, True, "fp8_block", "cuda:0") + assert third is not first, "one entry at a time" + + +def test_msa_layer_cache_keys_on_model_heads_gemm_and_device(): + cached, builds, evictions = _load( + "collect_msa_module.py", "_cached_msa_attention_layer", "create_msa_attention_layer" + ) + a = cached("MiniMaxAI/MiniMax-M3", 64, "bfloat16", "cuda:0") + assert cached("MiniMaxAI/MiniMax-M3", 64, "bfloat16", "cuda:0") is a and len(builds) == 1 + expected = {"model_path": "MiniMaxAI/MiniMax-M3", "num_heads": 64, "gemm_type": "bfloat16", "device": "cuda:0"} + assert builds[0] == expected + cached("MiniMaxAI/MiniMax-M3", 64, "fp8_block", "cuda:0") + assert len(builds) == 2 and evictions == ["gc", "empty_cache"] diff --git a/python/aisimulate/tests/unit/collector/trtllm/test_collect_moe_ep.py b/python/aisimulate/tests/unit/collector/trtllm/test_collect_moe_ep.py index fd6e800818..bbc60c27dc 100644 --- a/python/aisimulate/tests/unit/collector/trtllm/test_collect_moe_ep.py +++ b/python/aisimulate/tests/unit/collector/trtllm/test_collect_moe_ep.py @@ -415,30 +415,28 @@ def test_moe_ep_alone_resolves_to_the_wideep_trtllm_pin(): assert runtime.version == "1.3.0rc20" -def test_same_pin_mixing_with_stock_trtllm_ops_is_accepted(): - # Unlike wideep_sglang (0.5.10 vs stock 0.5.14 — mixing fail-closes), - # wideep_trtllm pins the SAME version AND image digest as stock trtllm - # (framework_manifest.yaml), so require_collector_runtime accepts a mixed - # stock+wideep request — the reason moe_ep can stay in the default trtllm - # model plans. The stock op here must live on the DEFAULT trtllm runtime: - # moe/sparse_attention carry per-family rc23 overrides, and mixing those - # with default-runtime ops correctly fail-closes as a multi-runtime - # request (a separate behavior from wideep/stock mixing). +def test_mixing_stock_and_wideep_trtllm_ops_fail_closes_on_different_pins(): + # Stock trtllm (1.3.0rc29 default) and wideep_trtllm (1.3.0rc20) no longer + # share a pin (framework_manifest.yaml, 2026-10-04): a mixed stock+wideep + # request is a multi-runtime request and must be split, exactly like + # wideep_sglang (0.5.10) vs stock sglang. moe_ep alone still resolves to + # the WideEP pin (test above); the stock op alone resolves to rc29. from collector.framework_manifest import require_collector_runtime - mixed = require_collector_runtime("trtllm", "1.3.0rc20", requested_ops={"gemm", "moe_ep"}, wideep_ops={"moe_ep"}) - stock = require_collector_runtime("trtllm", "1.3.0rc20", requested_ops={"gemm"}, wideep_ops={"moe_ep"}) - assert mixed.version == stock.version == "1.3.0rc20" - assert mixed.images == stock.images + with pytest.raises(RuntimeError, match=r"different runtime versions \(1\.3\.0rc29 != 1\.3\.0rc20\)"): + require_collector_runtime("trtllm", "1.3.0rc29", requested_ops={"gemm", "moe_ep"}, wideep_ops={"moe_ep"}) + stock = require_collector_runtime("trtllm", "1.3.0rc29", requested_ops={"gemm"}, wideep_ops={"moe_ep"}) + assert stock.framework == "trtllm" and stock.version == "1.3.0rc29" -def test_moe_ep_is_in_the_default_trtllm_plan_for_wideep_models(): +def test_moe_ep_is_requested_independently_of_the_default_trtllm_plan(): + # Mirrors sglang: with separate WideEP pins, moe_ep stays out of the default + # model plans on both backends so a default run resolves to one runtime. from collector.model_cases import build_collection_case_plan - trtllm_plan = build_collection_case_plan(backend="trtllm", model_path="deepseek-ai/DeepSeek-V3") - assert "moe_ep" in trtllm_plan.selected_ops - assert "trtllm_moe_wideep" not in trtllm_plan.selected_ops - # sglang keeps it out of the default plan (separate 0.5.10 runtime). + plan = build_collection_case_plan(backend="trtllm", model_path="deepseek-ai/DeepSeek-V3") + assert "moe_ep" not in plan.selected_ops + assert "trtllm_moe_wideep" not in plan.selected_ops sglang_plan = build_collection_case_plan(backend="sglang", model_path="deepseek-ai/DeepSeek-V3") assert "moe_ep" not in sglang_plan.selected_ops diff --git a/python/aisimulate/tests/unit/collector/trtllm/test_collect_moe_worker_recycle.py b/python/aisimulate/tests/unit/collector/trtllm/test_collect_moe_worker_recycle.py new file mode 100644 index 0000000000..6d546316dc --- /dev/null +++ b/python/aisimulate/tests/unit/collector/trtllm/test_collect_moe_worker_recycle.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""GPU-free tests for the TRT-LLM MoE worker-recycling policy. + +``collector/trtllm/collect_moe.py`` imports torch and tensorrt_llm at module +scope, so the policy helpers and the ``run_moe_torch`` wrapper are AST-extracted +and executed against fakes (same technique as ``test_collect_moe_ep.py``). +""" + +from __future__ import annotations + +import ast +import os +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +pytestmark = pytest.mark.unit + +REPO_ROOT = Path(__file__).resolve().parents[4] +SOURCE_PATH = REPO_ROOT / "collector" / "trtllm" / "collect_moe.py" +NAMES = ("_cases_since_worker_start", "_worker_restart_policy", "_recycle_worker_if_due", "run_moe_torch") +RELEASE = "_release_trtllm_memory_buffers" +EXIT_CODE_RESTART = 10 + +CASE_ARGS = ("nvfp4", [1, 2, 4], 7168, 2048, 8, 256, 1, 8, False, "deepseek-ai/DeepSeek-V3") +CASE_KWARGS = {"distributed": "power_law", "power_law_alpha": 1.2, "perf_filename": "moe_perf.txt", "device": "cuda:3"} + + +def _load(*, allocated_mb: float = 0.0): + tree = ast.parse(SOURCE_PATH.read_text(), filename=str(SOURCE_PATH)) + + def wanted(node) -> bool: + if isinstance(node, ast.FunctionDef): + return node.name in NAMES + if isinstance(node, ast.Assign): + return any(isinstance(target, ast.Name) and target.id in NAMES for target in node.targets) + return False + + selected = [node for node in tree.body if wanted(node)] + assert len(selected) == len(NAMES) + cases: list = [] + gc_calls: list = [] + namespace = { + "os": os, + "sys": sys, + "EXIT_CODE_RESTART": EXIT_CODE_RESTART, + "torch": SimpleNamespace( + device=lambda value: f"device:{value}", + cuda=SimpleNamespace( + memory_allocated=lambda device: allocated_mb * 2**20, + memory_reserved=lambda device: allocated_mb * 2**20, + ), + ), + "_print_live_cuda_tensors": lambda device: None, + "_release_trtllm_memory_buffers": lambda: 0, + "gc_collect": lambda: gc_calls.append(True), + "_run_moe_case": lambda *args, **kwargs: cases.append((args, kwargs)), + } + exec(compile(ast.Module(body=selected, type_ignores=[]), str(SOURCE_PATH), "exec"), namespace) + return namespace, cases, gc_calls + + +def _run_cases(namespace, count: int) -> None: + for _ in range(count): + namespace["run_moe_torch"](*CASE_ARGS, **CASE_KWARGS) + + +@pytest.fixture +def clean_env(monkeypatch): + for name in ("TRTLLM_MOE_RESTART_WORKER", "TRTLLM_MOE_RESTART_EVERY", "TRTLLM_MOE_RESTART_LEAK_MB"): + monkeypatch.delenv(name, raising=False) + + +def test_default_policy_recycles_after_sixteen_cases(clean_env): + namespace, cases, gc_calls = _load() + _run_cases(namespace, 15) + assert len(cases) == 15 and len(gc_calls) == 15 + with pytest.raises(SystemExit) as caught: + _run_cases(namespace, 1) + assert caught.value.code == EXIT_CODE_RESTART + assert len(cases) == 16, "the 16th case must still run (and persist) before the worker exits" + + +def test_leak_guard_recycles_immediately_when_memory_stays_allocated(clean_env): + namespace, cases, _ = _load(allocated_mb=2048) + with pytest.raises(SystemExit) as caught: + _run_cases(namespace, 1) + assert caught.value.code == EXIT_CODE_RESTART + assert len(cases) == 1 + + +def test_zero_disables_the_count_and_the_guard_and_the_legacy_switch_disables_recycling(clean_env, monkeypatch): + monkeypatch.setenv("TRTLLM_MOE_RESTART_EVERY", "0") + monkeypatch.setenv("TRTLLM_MOE_RESTART_LEAK_MB", "0") + namespace, cases, gc_calls = _load(allocated_mb=8192) + _run_cases(namespace, 50) + assert len(cases) == 50 and len(gc_calls) == 50 + # the legacy master switch wins over any count and over the default leak guard + monkeypatch.delenv("TRTLLM_MOE_RESTART_LEAK_MB") + monkeypatch.setenv("TRTLLM_MOE_RESTART_WORKER", "0") + monkeypatch.setenv("TRTLLM_MOE_RESTART_EVERY", "1") + namespace, cases, _ = _load(allocated_mb=8192) + _run_cases(namespace, 40) + assert len(cases) == 40 + + +def test_negative_values_are_rejected(clean_env, monkeypatch): + monkeypatch.setenv("TRTLLM_MOE_RESTART_EVERY", "-1") + namespace, _, _ = _load() + with pytest.raises(ValueError, match="must be >= 0"): + _run_cases(namespace, 1) + + +def test_wrapper_forwards_the_case_unchanged_and_recycles_after_it(clean_env, monkeypatch): + monkeypatch.setenv("TRTLLM_MOE_RESTART_EVERY", "2") + namespace, cases, gc_calls = _load() + _run_cases(namespace, 1) + assert cases == [(CASE_ARGS, CASE_KWARGS)] + assert gc_calls == [True], "cleanup runs after the case body returned" + + +def test_release_drops_every_trtllm_workspace_block_including_reserved(monkeypatch): + import types + + tree = ast.parse(SOURCE_PATH.read_text(), filename=str(SOURCE_PATH)) + node = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == RELEASE) + namespace: dict = {} + exec(compile(ast.Module(body=[node], type_ignores=[]), str(SOURCE_PATH), "exec"), namespace) + + class _Tensor: + def __init__(self, mib): + self.mib = mib + + def numel(self): + return self.mib * 2**20 + + def element_size(self): + return 1 + + class _Block: + def __init__(self, mib, reserved): + self.buffer = _Tensor(mib) + self.is_reserved = reserved + + blocks = {"moe_workspace": [_Block(7168, True), _Block(3584, False)], "other": [_Block(28, False)]} + pool = SimpleNamespace(buffers=blocks) + utils = types.ModuleType("tensorrt_llm._torch.memory_buffer_utils") + utils.get_memory_buffers = lambda: pool + pkg = types.ModuleType("tensorrt_llm._torch") + pkg.memory_buffer_utils = utils + monkeypatch.setitem(sys.modules, "tensorrt_llm", types.ModuleType("tensorrt_llm")) + monkeypatch.setitem(sys.modules, "tensorrt_llm._torch", pkg) + monkeypatch.setitem(sys.modules, "tensorrt_llm._torch.memory_buffer_utils", utils) + assert namespace[RELEASE]() == 7168 + 3584 + 28 + assert pool.buffers == {} diff --git a/python/aisimulate/tests/unit/collector/trtllm/test_mla_precision_mxfp4_padding.py b/python/aisimulate/tests/unit/collector/trtllm/test_mla_precision_mxfp4_padding.py index 01b4f66440..2f1968679e 100644 --- a/python/aisimulate/tests/unit/collector/trtllm/test_mla_precision_mxfp4_padding.py +++ b/python/aisimulate/tests/unit/collector/trtllm/test_mla_precision_mxfp4_padding.py @@ -56,7 +56,7 @@ class ReachedNativeBuilder(Exception): @pytest.mark.parametrize("model_name", ["openai/gpt-oss-120b", "openai/gpt-oss-20b", "other-mxfp4-model"]) def test_unaligned_mxfp4_honors_native_padding_window(sm, tp, quant, model_name): tree = ast.parse((COLLECTOR / "collect_moe.py").read_text()) - function = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "run_moe_torch") + function = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "_run_moe_case") native_builder = MagicMock(side_effect=ReachedNativeBuilder) namespace = { "torch": MagicMock(), @@ -72,17 +72,21 @@ def test_unaligned_mxfp4_honors_native_padding_window(sm, tp, quant, model_name) "get_sm_version": lambda: sm, "_MXFP4_MOE_TYPES": {"w4a16_mxfp4", "w4a8_mxfp4_mxfp8"}, "RenormalizeMoeRoutingMethod": MagicMock(), + # Declared-routing lookup (DeepSeek-V3-family gate) is not exercised by + # these MXFP4 checkpoints; keep the softmax-top-k path under test. + "_declared_moe_routing": lambda _model_name: None, + "DeepseekV3Gate": MagicMock(), "create_moe": native_builder, "inspect": SimpleNamespace(signature=lambda _: SimpleNamespace(parameters={})), } exec(compile(ast.Module(body=[function], type_ignores=[]), "collect_moe.py", "exec"), namespace) if sm in (90, 120): with pytest.raises(ValueError, match="weight-layout alignment"): - namespace["run_moe_torch"](quant, [1], 2880, 2880, 4, 128, tp, 1, False, model_name, perf_filename="unused") + namespace["_run_moe_case"](quant, [1], 2880, 2880, 4, 128, tp, 1, False, model_name, perf_filename="unused") native_builder.assert_not_called() return with pytest.raises(ReachedNativeBuilder): - namespace["run_moe_torch"](quant, [1], 2880, 2880, 4, 128, tp, 1, False, model_name, perf_filename="unused") + namespace["_run_moe_case"](quant, [1], 2880, 2880, 4, 128, tp, 1, False, model_name, perf_filename="unused") assert native_builder.call_args.kwargs["intermediate_size"] == 2880 assert native_builder.call_args.kwargs["model_config"].mapping.moe_tp_size == tp assert native_builder.call_args.kwargs["bias"] == model_name.startswith("openai/gpt-oss-") @@ -95,7 +99,7 @@ def test_unaligned_mxfp4_honors_native_padding_window(sm, tp, quant, model_name) @pytest.mark.parametrize("outcomes", ["fail,fail", "fail,pass", "pass,fail", "pass,pass", "fatal", "cached"]) def test_moe_autotuning_requires_success_or_loaded_cache(outcomes): tree = ast.parse((COLLECTOR / "collect_moe.py").read_text()) - function = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "run_moe_torch") + function = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "_run_moe_case") tuning = next( node for node in ast.walk(function) diff --git a/python/aisimulate/tests/unit/collector/trtllm/test_moe_serving_gate.py b/python/aisimulate/tests/unit/collector/trtllm/test_moe_serving_gate.py new file mode 100644 index 0000000000..955c13625b --- /dev/null +++ b/python/aisimulate/tests/unit/collector/trtllm/test_moe_serving_gate.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""GPU-free: the TRT-LLM MoE collector routes through the architecture's serving gate with a zeroed bias.""" + +import ast +import sys +from contextlib import nullcontext +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + +pytestmark = pytest.mark.unit +COLLECT_MOE = Path(__file__).resolve().parents[4] / "collector" / "trtllm" / "collect_moe.py" + + +class _Bias: + zeroed = False + + def zero_(self): + self.zeroed = True + + +def _gate(kind): + class Gate: + def __init__(self, *args, **kwargs): + self.kind = kind + self.kwargs = kwargs + self.e_score_correction_bias = _Bias() + + return Gate + + +@pytest.fixture +def serving_gate(monkeypatch): + tree = ast.parse(COLLECT_MOE.read_text()) + helpers = [ + n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name in ("_zero_routing_bias", "_serving_moe_gate") + ] + namespace = {"torch": SimpleNamespace(no_grad=nullcontext), "DeepseekV3Gate": _gate("v3")} + exec(compile(ast.Module(body=helpers, type_ignores=[]), str(COLLECT_MOE), "exec"), namespace) + for name in ("tensorrt_llm", "tensorrt_llm._torch", "tensorrt_llm._torch.models"): + monkeypatch.setitem(sys.modules, name, ModuleType(name)) + for name, attr, kind in ( + ("modeling_deepseekv4", "DeepseekV4Gate", "v4"), + ("modeling_minimaxm3", "MiniMaxM3Gate", "m3"), + ): + module = ModuleType(f"tensorrt_llm._torch.models.{name}") + setattr(module, attr, _gate(kind)) + monkeypatch.setitem(sys.modules, module.__name__, module) + return namespace["_serving_moe_gate"] + + +@pytest.mark.parametrize( + ("architecture", "kind"), + [ + ("DeepseekV3ForCausalLM", "v3"), + ("GlmMoeDsaForCausalLM", "v3"), + ("DeepseekV4ForCausalLM", "v4"), + ("MiniMaxM3ForCausalLM", "m3"), + ], +) +def test_gate_matches_the_architecture_and_its_bias_is_zeroed(serving_gate, architecture, kind): + declared = SimpleNamespace( + architecture=architecture, + sglang_moe_num_expert_group=8, + sglang_moe_topk_group=4, + sglang_moe_routed_scaling_factor=2.5, + ) + gate = serving_gate(declared, 7168, 256, 8, "bf16", "trtllm") + assert gate.kind == kind + assert gate.e_score_correction_bias.zeroed + if kind == "v3": + # serving compares the resolved backend upper-case: 'TRTLLM' keeps a bf16 bias + assert gate.kwargs["moe_backend"] == "TRTLLM" + if kind == "v4": + assert gate.kwargs["is_hashed"] is False diff --git a/python/aisimulate/tests/unit/collector/vllm/test_collect_kda_contract.py b/python/aisimulate/tests/unit/collector/vllm/test_collect_kda_contract.py index 9250f6bb21..7762971329 100644 --- a/python/aisimulate/tests/unit/collector/vllm/test_collect_kda_contract.py +++ b/python/aisimulate/tests/unit/collector/vllm/test_collect_kda_contract.py @@ -267,25 +267,6 @@ def test_glm5_next_decode_mirrors_serving_decode(): assert ast.unparse(keywords["lower_bound"]) == "GLM5_NEXT_KDA_LOWER_BOUND" -@pytest.mark.parametrize( - ("version", "accepted"), - [ - ("0.1.dev19262+gb6bbf29dd", True), - ("0.31.0", True), - ("0.31.1", False), - ], -) -def test_kda_compat_admits_kimi_preview_and_glm_runtime(version, accepted): - from collector.version_resolver import _check_compat - - declaration = next( - node.value.value - for node in ast.parse(SOURCE_PATH.read_text(encoding="utf-8")).body - if isinstance(node, ast.Assign) and any(getattr(t, "id", None) == "__compat__" for t in node.targets) - ) - assert _check_compat(declaration, version) is accepted - - _GATE_NAMES = { "GLM5_NEXT_KDA_MODEL_PATHS", "_is_glm5_next_kda", @@ -313,23 +294,13 @@ def _gate_namespace(): return namespace -def test_kda_audited_release_sets_do_not_overlap(): - from collector.version_resolver import _check_compat - - compat = _gate_namespace()["_KDA_ARCHITECTURE_COMPAT"] - audited = {arch: spec.split("==", 1)[1] for arch, spec in compat.items()} - assert set(audited) == {"KimiK3ForConditionalGeneration", "Glm5NextForConditionalGeneration"} - for arch, version in audited.items(): - for other, spec in compat.items(): - assert _check_compat(spec, version) is (other == arch), (arch, other) - - @pytest.mark.parametrize( ("model_name", "version", "audited"), [ - ("moonshotai/Kimi-K3", "0.1.dev19262+gb6bbf29dd", True), + ("moonshotai/Kimi-K3", "0.31.0", True), + ("moonshotai/Kimi-K3", "0.1.dev19262+gb6bbf29dd", False), ("moonshotai/Kimi-K3", "0.27.0", False), - ("moonshotai/Kimi-K3", "0.31.0", False), + ("moonshotai/Kimi-K3", "0.30.0", False), ("zai-org/GLM-5.3-Flash", "0.31.0", True), ("nvidia/GLM-5.3-Flash-NVFP4", "0.31.0", True), ("zai-org/GLM-5.3-Flash", "0.30.0", False), diff --git a/python/aisimulate/tests/unit/collector/vllm/test_rope_missing_key_shim.py b/python/aisimulate/tests/unit/collector/vllm/test_rope_missing_key_shim.py new file mode 100644 index 0000000000..e49a154cb6 --- /dev/null +++ b/python/aisimulate/tests/unit/collector/vllm/test_rope_missing_key_shim.py @@ -0,0 +1,246 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""vLLM 0.30/0.31 rope q_pe alone (key=None) in the DSV3.2 dense-MHA FP8 branch +(models/deepseek_v32/attention.py:534 @0.30.0). Only DeepseekScalingRotaryEmbedding +rejects that; the collector wraps THAT class's forward so the call ropes the query +against a throwaway key that shares the query's strides and returns (query, None). +Rotary classes that accept key=None (GLM-5.x's RotaryEmbedding) stay on serving's call: +wrapping them handed vLLM's own rope kernel — one head stride for query and key — a +contiguous key with the wrong head stride, i.e. out-of-bounds in-place writes (the +GLM-5.2 fp8-KV dsa_context illegal-address / stall family, 2026-10-14).""" + +import ast +import math +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.unit +SOURCE = Path(__file__).resolve().parents[4] / "collector" / "vllm" / "collect_mla_module.py" +_FUNCTIONS = ( + "_strided_throwaway_copy", + "_rope_tolerating_missing_key", + "_rope_module_rejects_missing_key", + "_install_rope_missing_key_shims", +) + + +def _load_all(): + tree = ast.parse(SOURCE.read_text(encoding="utf-8"), filename=str(SOURCE)) + fns = [n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name in _FUNCTIONS] + assert {f.name for f in fns} == set(_FUNCTIONS) + ns: dict = {} + exec(compile(ast.Module(body=fns, type_ignores=[]), str(SOURCE), "exec"), ns) + return ns + + +def _load(): + return _load_all()["_rope_tolerating_missing_key"] + + +def _contiguous_strides(shape): + strides = [] + acc = 1 + for s in reversed(shape): + strides.append(acc) + acc *= s + return tuple(reversed(strides)) + + +class _Q: + """Just enough of torch.Tensor for the shim: shape/stride and the four methods it calls.""" + + def __init__(self, tag, shape=(4, 8, 64), stride=None, storage_numel=None): + self.tag = tag + self.shape = tuple(shape) + self._stride = tuple(stride) if stride is not None else _contiguous_strides(self.shape) + self.storage_numel = storage_numel if storage_numel is not None else math.prod(self.shape) + self.made_by = "fresh" + + def stride(self): + return self._stride + + def clone(self): + q = _Q(self.tag + "'", self.shape) + q.made_by = "clone" + return q + + def new_empty(self, numel): + return _Q("empty", (numel,), storage_numel=numel) + + def as_strided(self, shape, stride): + q = _Q("view", shape, stride, storage_numel=self.storage_numel) + q.made_by = "as_strided" + return q + + def copy_(self, src): + self.tag = src.tag + "'" + return self + + +def _q_pe_slice(tokens=3, heads=8, qk_head_dim=256, rope_dim=64): + # q_pe = q.split([nope, rope], dim=-1)[1]: a last-dim slice of [tokens, heads, qk_head_dim]. + return _Q("q", (tokens, heads, rope_dim), (heads * qk_head_dim, qk_head_dim, 1)) + + +def test_missing_key_is_replaced_by_a_throwaway_copy_and_dropped(): + calls = [] + + def forward(positions, query, key=None, offsets=None): + assert key is not None # the real impls' contract + calls.append((positions, query.tag, key.tag, offsets)) + return (_Q(query.tag + "+rope"), _Q(key.tag + "+rope")) + + wrapped = _load()(forward) + q, k = wrapped("pos", _Q("q")) + assert (q.tag, k) == ("q+rope", None) + assert calls == [("pos", "q", "q'", None)] + + +def test_throwaway_key_shares_the_query_strides_not_a_contiguous_clone(): + # vLLM's rope kernel addresses the key with the QUERY's head stride + # (csrc/libtorch_stable/pos_encoding_kernels.cu:57,68,164-165 @0.30.0). A contiguous clone of + # the q_pe slice (head stride 64 vs 256) is written past its end by 192*(heads-1) elements. + seen = [] + + def forward(positions, query, key=None): + seen.append(key) + return (query, key) + + query = _q_pe_slice(tokens=3, heads=8) + _load()(forward)("pos", query) + (key,) = seen + assert key.made_by == "as_strided" + assert key.shape == query.shape + assert key.stride() == query.stride() == (8 * 256, 256, 1) + # Storage covers the last addressable element of the strided layout, nothing less. + assert key.storage_numel == 1 + 2 * (8 * 256) + 7 * 256 + 63 + assert key.tag == "q'" # the data was copied from the query + + +def test_present_key_passes_through_unchanged(): + def forward(positions, query, key=None, offsets=None): + return (query, key) + + wrapped = _load()(forward) + assert wrapped("pos", "q", "k", "off") == ("q", "k") + + +def test_shim_forwards_only_the_arguments_it_received(): + # vLLM 0.30 forward_cuda(positions, query, key=None): no offsets parameter. The first shim + # passed offsets=None positionally -> TypeError "takes from 3 to 4 positional arguments but + # 5 were given" (508 GLM-5 cases, pipelines 72387095/72387091). + def forward_cuda(positions, query, key=None): + return (query, key) + + wrapped = _load()(forward_cuda) + assert wrapped("pos", "q", "k") == ("q", "k") + q, k = wrapped("pos", _Q("q")) + assert (q.tag, k) == ("q", None) + + +def test_shim_is_idempotent(): + # Installed per case on a possibly reused module: re-wrapping stacked closures until + # RecursionError (2,286 cases per run at e598d380). + def forward(positions, query, key=None): + return (query, key) + + shim = _load() + once = shim(forward) + assert shim(once) is once and shim(shim(once)) is once + assert getattr(once, "_aisim_rope_shim", False) is True + + # ... and the installer, re-run on the same module for every case, keeps that one wrapper. + install = _load_all()["_install_rope_missing_key_shims"] + dsv32_rope = DeepseekScalingRotaryEmbedding() + attn = _Attn(rotary_emb=dsv32_rope) + install(attn) + wrapped = dsv32_rope.forward + install(attn) + assert dsv32_rope.forward is wrapped + + +# --- which rotary classes get the shim ------------------------------------------------- + + +class RotaryEmbedding: + """vLLM's base class: forward_cuda -> ops.rotary_embedding, key may be None.""" + + def forward_cuda(self, positions, query, key=None): + return (query, key) + + def forward(self, positions, query, key=None): + return self.forward_cuda(positions, query, key) + + +class Llama3RotaryEmbedding(RotaryEmbedding): + pass + + +class DeepseekScalingRotaryEmbedding(RotaryEmbedding): + """DeepSeek-V3.2's yarn rope: FlashInfer op key.view / native assert key is not None.""" + + def forward_cuda(self, positions, query, key=None): + assert key is not None, "the real impls reject key=None" + return (query, key) + + +class DeepseekV4ScalingRotaryEmbedding(DeepseekScalingRotaryEmbedding): + """Overrides forward_cuda back onto ops.rotary_embedding: accepts key=None again.""" + + def forward_cuda(self, positions, query, key=None): + return (query, key) + + +class Linear: + def forward(self, x): + return x + + +class _Attn: + def __init__(self, **modules): + self._modules = modules + + def named_modules(self): + return list(self._modules.items()) + + +def test_rejects_missing_key_only_for_deepseek_scaling_rope_itself(): + rejects = _load_all()["_rope_module_rejects_missing_key"] + assert rejects(DeepseekScalingRotaryEmbedding()) is True + assert rejects(RotaryEmbedding()) is False # GLM-5.x: rope_type "default" + assert rejects(Llama3RotaryEmbedding()) is False + assert rejects(DeepseekV4ScalingRotaryEmbedding()) is False + assert rejects(Linear()) is False + + +def test_installer_shims_deepseek_scaling_rope_and_leaves_glm_rope_on_servings_call(): + ns = _load_all() + glm_rope = RotaryEmbedding() + dsv32_rope = DeepseekScalingRotaryEmbedding() + attn = _Attn( + **{ + "rotary_emb": glm_rope, + "indexer_rope_emb": Llama3RotaryEmbedding(), + "mla_attn.rotary_emb": dsv32_rope, + "q_b_proj": Linear(), + } + ) + patched, left = ns["_install_rope_missing_key_shims"](attn) + assert patched == ["DeepseekScalingRotaryEmbedding"] + assert left == ["RotaryEmbedding", "Llama3RotaryEmbedding"] + + # GLM's rope is untouched: serving's `rotary_emb(positions, q_pe)` reaches the class's own + # forward with key=None, exactly as in vLLM. + assert "forward" not in vars(glm_rope) + assert glm_rope.forward("pos", "q") == ("q", None) + + # DeepSeek-V3.2's rope now accepts the same call and returns (query, None). + assert getattr(dsv32_rope.forward, "_aisim_rope_shim", False) is True + q, k = dsv32_rope.forward("pos", _q_pe_slice()) + assert (q.tag, k) == ("q", None) + + +def test_run_mla_module_installs_shims_through_the_class_filter(): + text = SOURCE.read_text(encoding="utf-8") + assert "_install_rope_missing_key_shims(attn_module)" in text diff --git a/python/aisimulate/tests/unit/generator/test_engine_limit_preservation.py b/python/aisimulate/tests/unit/generator/test_engine_limit_preservation.py index 09464f3830..ec4521ce65 100644 --- a/python/aisimulate/tests/unit/generator/test_engine_limit_preservation.py +++ b/python/aisimulate/tests/unit/generator/test_engine_limit_preservation.py @@ -77,10 +77,22 @@ def test_externally_evaluated_engine_limits_are_preserved(backend: str, rule: st assert actual[key] == expected[key] +@pytest.mark.parametrize("backend", ["vllm", "sglang", "trtllm"]) +def test_preserved_unaligned_budget_is_rounded_with_a_warning(backend, caplog): + """A Sweeper candidate simulated at 5000 tokens deploys 5056; the change is logged, not silent.""" + params = _params(preserve=True, rule="benchmark") + params["params"]["prefill"]["max_num_tokens"] = 5000 + params["params"]["prefill"]["max_prefill_tokens"] = 5000 + with caplog.at_level(logging.WARNING): + result = apply_rule_plugins(copy.deepcopy(params), backend) + assert result["params"]["prefill"]["max_num_tokens"] == 5056 + assert "prefill max_num_tokens=5000 was evaluated by the caller; the deployment uses 5056" in caplog.text + + def test_default_vllm_rules_still_resize_unpinned_requests(): result = apply_rule_plugins(_params(preserve=False), "vllm") - assert result["params"]["prefill"]["max_num_tokens"] == 1628 + assert result["params"]["prefill"]["max_num_tokens"] == 1664 # 128 + 1500 = 1628, aligned up to 64 assert result["params"]["decode"]["max_batch_size"] == 512 assert result["params"]["decode"]["max_num_tokens"] == 512 diff --git a/python/aisimulate/tests/unit/generator/test_facts_hardware_l40s.py b/python/aisimulate/tests/unit/generator/test_facts_hardware_l40s.py new file mode 100644 index 0000000000..c75f8b5837 --- /dev/null +++ b/python/aisimulate/tests/unit/generator/test_facts_hardware_l40s.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ada (sm89) hardware facts: the l40s profile, and the quantization-keyed trtllm MoE backend. + +Without the profile resolve_facts() raised "Unknown hardware profile 'l40s'" and the pipeline swallowed +it (no model default reached l40s renders). On SM89 trtllm 1.3.0rc29 serves NVFP4 MoE only through +MARLIN, which rejects bf16/fp8 MoE layers, so the profile keys the choice by quantization family +(op-probe harness L40 campaign, 2026-10-04, findings sm89_trtllm_rc29_probe_2026_10_04).""" + +import pytest + +from aisimulate.generator.facts.apply import apply_moe_backend, quant_family_of +from aisimulate.generator.facts.request_resolution import hardware_key_for_system +from aisimulate.generator.facts.resolve import _FACTS_DIR, load_backend_version_matrix, resolve_facts + +pytestmark = pytest.mark.unit + + +def _latest_dynamo() -> str: + matrix = load_backend_version_matrix(str(_FACTS_DIR / "runtimes" / "dynamo.yaml")) + return next(v for v, backends in matrix.items() if "trtllm" in backends) + + +def _l40s(): + return resolve_facts( + model_profile_id=None, hardware="l40s", transport="nvlink", dynamo_version=_latest_dynamo(), backend="trtllm" + ).hardware + + +def test_l40s_is_its_own_profile_and_no_phantom_l40_system(): + assert hardware_key_for_system("l40s") == "l40s" + hw = _l40s() + assert hw["node_selector"]["nvidia.com/gpu.product"] == "NVIDIA-L40S" + + +@pytest.mark.parametrize( + "quant, family", + [ + (None, None), + ({"quant_algo": "NVFP4", "kv_cache_quant_algo": "FP8"}, "nvfp4"), # modelopt hf_quant_config + ({"quant_method": "mxfp4"}, "mxfp4"), # gpt-oss + ({"quant_method": "fp8", "weight_block_size": [128, 128]}, "fp8_block"), # DeepSeek-V3 class + ({"quant_method": "fp8"}, "fp8"), + # modelopt MIXED_PRECISION: the experts' algo decides (nvidia/Qwen3.8-2.4T-A95B-NVFP4 layout) + ( + { + "quant_algo": "MIXED_PRECISION", + "quant_method": "modelopt", + "quantized_layers": { + "model.layers.0.linear_attn.in_proj_qkv": {"quant_algo": "FP8"}, + "model.layers.0.linear_attn.out_proj": {"quant_algo": "FP8"}, + "model.layers.0.mlp.experts": {"quant_algo": "NVFP4"}, + }, + }, + "nvfp4", + ), + ( + { + "quant_algo": "MIXED_PRECISION", + "quant_method": "modelopt", + "quantized_layers": { + "model.layers.0.self_attn.q_proj": {"quant_algo": "FP8"}, + "model.layers.0.mlp.down_proj": {"quant_algo": "FP8"}, + }, + }, + "fp8", + ), + ], +) +def test_quant_family_of(quant, family): + assert quant_family_of(quant) == family + + +def test_l40s_fills_marlin_for_nvfp4_moe_only(): + hw = _l40s() + ctx = {"is_moe": True, "moe_config": {}} + apply_moe_backend(ctx, hw, backend="trtllm", quantization={"quant_algo": "NVFP4"}) + assert ctx["moe_config"]["backend"] == "MARLIN" + for quant in (None, {"quant_method": "fp8", "weight_block_size": [128, 128]}, {"quant_method": "mxfp4"}): + ctx = {"is_moe": True, "moe_config": {}} + apply_moe_backend(ctx, hw, backend="trtllm", quantization=quant) + assert "backend" not in ctx["moe_config"], quant # the template default applies (rc29: AUTO -> CUTLASS on SM89) + dense = {"is_moe": False, "moe_config": {}} + apply_moe_backend(dense, hw, backend="trtllm", quantization={"quant_algo": "NVFP4"}) + assert "backend" not in dense["moe_config"] + + +def test_user_moe_backend_param_beats_the_hardware_fill_on_trtllm(): + """`--generator-set params.agg.moe_backend=...` used to be silently dropped on trtllm: the hardware fact + filled moe_config.backend regardless and the template never read the generic param.""" + hw = {"moe_backend": {"trtllm": "CUTLASS"}} + ctx = {"is_moe": True, "moe_config": {}, "agg_moe_backend": "marlin", "moe_backend": "marlin"} + apply_moe_backend(ctx, hw, backend="trtllm") + assert "backend" not in ctx["moe_config"] + ctx = {"is_moe": True, "moe_config": {}} + apply_moe_backend(ctx, hw, backend="trtllm") + assert ctx["moe_config"]["backend"] == "CUTLASS" diff --git a/python/aisimulate/tests/unit/generator/test_facts_hardware_sm120.py b/python/aisimulate/tests/unit/generator/test_facts_hardware_sm120.py new file mode 100644 index 0000000000..766c232076 --- /dev/null +++ b/python/aisimulate/tests/unit/generator/test_facts_hardware_sm120.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""RTX PRO (sm120) hardware facts profile. Without it resolve_facts() raised +"Unknown hardware profile 'rtx_pro_6000_server'", the pipeline swallowed the +error and every sm120 render lost its model defaults (op-probe harness, +2026-10-01: V4.1's block-size 64 never reached sm120 engine args).""" + +import pytest + +from aisimulate.generator.facts.apply import apply_model_default_args +from aisimulate.generator.facts.request_resolution import hardware_key_for_system +from aisimulate.generator.facts.resolve import _FACTS_DIR, load_backend_version_matrix, resolve_facts + +pytestmark = pytest.mark.unit + + +def _latest_dynamo() -> str: + matrix = load_backend_version_matrix(str(_FACTS_DIR / "runtimes" / "dynamo.yaml")) + return next(v for v, backends in matrix.items() if "vllm" in backends) + + +def test_rtx_pro_6000_server_maps_to_the_sm120_profile(): + assert hardware_key_for_system("rtx_pro_6000_server") == "sm120" + # unknown systems still fall through to the raw name (clear error downstream) + assert hardware_key_for_system("not_a_system") == "not_a_system" + + +def test_sm120_profile_resolves_and_model_defaults_reach_its_engine_args(): + facts = resolve_facts( + model_profile_id="deepseek-v4.1", + hardware="sm120", + transport="nvlink", + dynamo_version=_latest_dynamo(), + backend="vllm", + ) + # PCIe Blackwell: no wide-EP MoE backend is asserted for any backend + assert "moe_backend" not in facts.hardware + tokens: list[str] = [] + apply_model_default_args(tokens, facts.model, backend="vllm", system="sm120", role="agg", variant=None) + assert tokens[tokens.index("--block-size") + 1] == "64" + assert "--trust-remote-code" in tokens diff --git a/python/aisimulate/tests/unit/generator/test_llmd_deployment.py b/python/aisimulate/tests/unit/generator/test_llmd_deployment.py index 017781f36a..d6a511883c 100644 --- a/python/aisimulate/tests/unit/generator/test_llmd_deployment.py +++ b/python/aisimulate/tests/unit/generator/test_llmd_deployment.py @@ -174,7 +174,8 @@ def test_llmd_kustomize_vllm_disagg_mode(): assert prefill["spec"]["replicas"] == 14 assert prefill_container["resources"]["limits"]["nvidia.com/gpu"] == "1" assert "--max-num-batched-tokens" in prefill_container["args"] - assert "5500" in prefill_container["args"] + # every token budget renders 64-aligned (owner decision 2026-09-24): 5500 -> 5504 + assert "5504" in prefill_container["args"] @pytest.mark.unit diff --git a/python/aisimulate/tests/unit/generator/test_model_profile_deepseek_v41.py b/python/aisimulate/tests/unit/generator/test_model_profile_deepseek_v41.py new file mode 100644 index 0000000000..337054daa6 --- /dev/null +++ b/python/aisimulate/tests/unit/generator/test_model_profile_deepseek_v41.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DeepSeek-V4.1 has its own facts profile: V4's `DeepSeek-V4*` glob used to +hand it block-size 256, which vllm 0.30.0 rejects for V4.1's 64-token kernel +blocks (op-probe harness A/B, 2026-10-01).""" + +import pytest + +from aisimulate.generator.facts.apply import apply_model_default_args +from aisimulate.generator.facts.request_resolution import model_profile_for_path +from aisimulate.generator.facts.resolve import resolve_facts + +pytestmark = pytest.mark.unit + + +@pytest.mark.parametrize( + "path,profile", + [ + ("deepseek-ai/DeepSeek-V4-Flash-0731", "deepseek-v4"), + ("deepseek-ai/DeepSeek-V4.1-Flash", "deepseek-v4.1"), + ("nvidia/DeepSeek-V4.1-Flash-NVFP4", "deepseek-v4.1"), + ], +) +def test_v4_and_v41_resolve_to_their_own_profiles(path, profile): + assert model_profile_for_path(path) == profile + + +def _vllm_defaults(profile_id: str) -> list[str]: + facts = resolve_facts( + model_profile_id=profile_id, + hardware="h200", + transport="nvlink", + dynamo_version=_latest_dynamo(), + backend="vllm", + ) + tokens: list[str] = [] + apply_model_default_args(tokens, facts.model, backend="vllm", system=None, role="agg", variant=None) + return tokens + + +def _latest_dynamo() -> str: + from aisimulate.generator.facts.resolve import _FACTS_DIR, load_backend_version_matrix + + matrix = load_backend_version_matrix(str(_FACTS_DIR / "runtimes" / "dynamo.yaml")) + return next(v for v, backends in matrix.items() if "vllm" in backends) + + +def test_v41_vllm_block_size_is_64_not_v4s_256(): + v41 = _vllm_defaults("deepseek-v4.1") + v4 = _vllm_defaults("deepseek-v4") + assert v41[v41.index("--block-size") + 1] == "64" + assert v4[v4.index("--block-size") + 1] == "256" + assert "--trust-remote-code" in v41 and "--no-enable-flashinfer-autotune" in v41 diff --git a/python/aisimulate/tests/unit/generator/test_model_profile_minimax_m3.py b/python/aisimulate/tests/unit/generator/test_model_profile_minimax_m3.py new file mode 100644 index 0000000000..afd734b6df --- /dev/null +++ b/python/aisimulate/tests/unit/generator/test_model_profile_minimax_m3.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 model facts: vllm needs block-size 128 because the sparse layers' backend accepts only that KV page +size while the dense head's FlashAttention accepts MultipleOf(16) (op-probe harness, L40, 2026-10-05).""" + +import pytest + +from aisimulate.generator.facts.apply import apply_model_default_args +from aisimulate.generator.facts.request_resolution import model_profile_for_path +from aisimulate.generator.facts.resolve import _FACTS_DIR, load_backend_version_matrix, resolve_facts + +pytestmark = pytest.mark.unit + + +def _latest_dynamo(backend: str) -> str: + matrix = load_backend_version_matrix(str(_FACTS_DIR / "runtimes" / "dynamo.yaml")) + return next(v for v, backends in matrix.items() if backend in backends) + + +@pytest.mark.parametrize("path", ["MiniMaxAI/MiniMax-M3", "MiniMaxAI/MiniMax-M3-MXFP8", "nvidia/MiniMax-M3-NVFP4"]) +def test_minimax_m3_paths_resolve_to_the_profile(path): + assert model_profile_for_path(path) == "minimax-m3" + + +def test_minimax_m3_vllm_default_is_block_size_128_only(): + def tokens_for(backend: str) -> list[str]: + facts = resolve_facts( + model_profile_id="minimax-m3", + hardware="h200", + transport="nvlink", + dynamo_version=_latest_dynamo(backend), + backend=backend, + ) + out: list[str] = [] + apply_model_default_args(out, facts.model, backend=backend, system=None, role="agg", variant=None) + return out + + assert tokens_for("vllm") == ["--block-size", "128"] + # sglang / trtllm carry no M3 default from this profile + assert tokens_for("sglang") == [] + assert tokens_for("trtllm") == [] diff --git a/python/aisimulate/tests/unit/generator/test_sglang_cuda_graph_flags_0520.py b/python/aisimulate/tests/unit/generator/test_sglang_cuda_graph_flags_0520.py new file mode 100644 index 0000000000..beb4332fdc --- /dev/null +++ b/python/aisimulate/tests/unit/generator/test_sglang_cuda_graph_flags_0520.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SGLang >= 0.5.20 renamed `--cuda-graph-bs` / `--cuda-graph-max-bs` to the +`*-decode` spellings (ServerArgs is a msgspec Struct; the old aliases are gone, +0.5.21 rejects `--cuda-graph-bs` as ambiguous). cli_args.0.5.20.j2 renders the +new spelling; older versions keep the old one (op-probe harness, 2026-10-01).""" + +import copy + +import pytest +import yaml + +from aisimulate.generator.api import generate_backend_artifacts + +pytestmark = pytest.mark.unit + +_PARAMS = { + "ServiceConfig": { + "model_path": "deepseek-ai/DeepSeek-V3", + "served_model_path": "deepseek-ai/DeepSeek-V3", + "served_model_name": "DeepSeek-V3", + "include_frontend": True, + }, + "K8sConfig": {"name_prefix": "test", "k8s_namespace": "default"}, + "DynConfig": {"mode": "agg"}, + "WorkerConfig": {"agg_workers": 1, "agg_gpus_per_worker": 8, "prefill_workers": 0, "decode_workers": 0}, + "NodeConfig": {"num_gpus_per_node": 8}, + "SlaConfig": {"isl": 1024, "osl": 256}, + "ModelConfig": {"is_moe": True, "prefix": 0, "nextn": 0}, + "BenchConfig": {}, + "params": { + "agg": { + "tensor_parallel_size": 8, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + "moe_tensor_parallel_size": 1, + "moe_expert_parallel_size": 8, + "max_batch_size": 64, + "max_num_tokens": 4096, + "max_seq_len": 4096, + } + }, +} + + +def _worker_args(version: str) -> str: + artifacts = generate_backend_artifacts( + copy.deepcopy(_PARAMS), "sglang", backend_version=version, deployment_target="dynamo-j2" + ) + k8s = yaml.safe_load(artifacts["k8s_deploy.yaml"]) + worker = next(svc for name, svc in k8s["spec"]["services"].items() if name != "Frontend") + return " ".join(worker["extraPodSpec"]["mainContainer"]["args"]) + + +@pytest.mark.parametrize("version", ["0.5.20", "0.5.21"]) +def test_new_versions_render_decode_spelling(version): + args = _worker_args(version) + assert "--cuda-graph-bs-decode " in args + assert "--cuda-graph-bs " not in args + assert "--cuda-graph-max-bs " not in args + + +@pytest.mark.parametrize("version", ["0.5.16", "0.5.19"]) +def test_old_versions_keep_old_spelling(version): + args = _worker_args(version) + assert "--cuda-graph-bs " in args + assert "--cuda-graph-bs-decode" not in args + + +def test_prefill_graph_backend_param_renders_only_on_new_versions(): + """params.agg.cuda_graph_backend_prefill -> --cuda-graph-backend-prefill (sglang >= 0.5.20 template).""" + params = copy.deepcopy(_PARAMS) + params["params"]["agg"]["cuda_graph_backend_prefill"] = "breakable" + new = generate_backend_artifacts(params, "sglang", backend_version="0.5.21", deployment_target="dynamo-j2") + k8s = yaml.safe_load(new["k8s_deploy.yaml"]) + worker = next(s for n, s in k8s["spec"]["services"].items() if n != "Frontend") + assert '--cuda-graph-backend-prefill "breakable"' in " ".join(worker["extraPodSpec"]["mainContainer"]["args"]) + old = generate_backend_artifacts( + copy.deepcopy(params), "sglang", backend_version="0.5.16", deployment_target="dynamo-j2" + ) + k8s = yaml.safe_load(old["k8s_deploy.yaml"]) + worker = next(s for n, s in k8s["spec"]["services"].items() if n != "Frontend") + assert "--cuda-graph-backend-prefill" not in " ".join(worker["extraPodSpec"]["mainContainer"]["args"]) diff --git a/python/aisimulate/tests/unit/generator/test_sweeper_request.py b/python/aisimulate/tests/unit/generator/test_sweeper_request.py index c6149efc44..0be267add1 100644 --- a/python/aisimulate/tests/unit/generator/test_sweeper_request.py +++ b/python/aisimulate/tests/unit/generator/test_sweeper_request.py @@ -815,3 +815,29 @@ def _assert_artifact_block_size(artifacts, backend, role, expected): assert flag not in shlex.split(artifacts[f"cli_args_{role}"]) else: assert _cli_flag_value(artifacts[f"cli_args_{role}"], flag) == str(expected) + + +@pytest.mark.parametrize(("hardware_sku", "implementation"), [("h200_sxm", None), ("b200_sxm", "msa")]) +def test_minimax_m3_trtllm_candidate_renders_its_sparse_attention_keys(hardware_sku, implementation): + """The Sweeper lowers through rule=benchmark; rc29 cannot build the M3 KV-cache manager + without sparse_attention_config.algorithm, on every SM.""" + candidate = _agg_candidate( + model_name="MiniMaxAI/MiniMax-M3", + backend="trtllm", + backend_version="1.3.0rc29", + hardware_sku=hardware_sku, + moe_ep=2, + adapters={}, + ) + request = from_sweeper_candidate( + candidate, + workload={"isl": 4000, "osl": 1000}, + model_facts=ModelFacts(is_moe=True, architecture="MiniMaxM3ForCausalLM"), + generator_overrides={"K8sConfig": {"k8s_image": "example/trtllm:1.3.0rc29"}}, + ) + artifacts = generate_from_request(request) + name = next(name for name in artifacts if "extra_engine_args" in name and "agg" in name) + expected = {"algorithm": "minimax_m3"} + if implementation: + expected["implementation"] = implementation + assert yaml.safe_load(artifacts[name]).get("sparse_attention_config") == expected diff --git a/python/aisimulate/tests/unit/generator/test_trtllm_moe_backend_default.py b/python/aisimulate/tests/unit/generator/test_trtllm_moe_backend_default.py new file mode 100644 index 0000000000..372bf74479 --- /dev/null +++ b/python/aisimulate/tests/unit/generator/test_trtllm_moe_backend_default.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""TRT-LLM MoE backend on Blackwell: WIDEEP only for releases whose MoeConfig accepts it +(through 1.3.0rc24); from rc29 an unset backend renders the framework default AUTO.""" + +from __future__ import annotations + +import pytest +import yaml + +from aisimulate.generator.api import generate_backend_artifacts +from aisimulate.generator.facts.apply import apply_moe_backend + +pytestmark = pytest.mark.unit + + +@pytest.mark.parametrize( + ("version", "expected"), + [("1.3.0rc14", "WIDEEP"), ("1.3.0rc24", "WIDEEP"), ("1.3.0rc25", None), ("1.3.0rc29", None), (None, "WIDEEP")], +) +def test_wideep_is_filled_only_for_releases_that_accept_it(version, expected): + ctx = {"is_moe": True, "moe_config": {}} + apply_moe_backend(ctx, {"moe_backend": {"trtllm": "WIDEEP"}}, backend="trtllm", backend_version=version) + assert ctx["moe_config"].get("backend") == expected + + +def _params(system: str) -> dict: + return { + "ServiceConfig": { + "model_path": "deepseek-ai/DeepSeek-V3", + "served_model_path": "deepseek-ai/DeepSeek-V3", + "served_model_name": "DeepSeek-V3", + "include_frontend": True, + }, + "K8sConfig": {"name_prefix": "test", "k8s_namespace": "default"}, + "DynConfig": {"mode": "agg"}, + "WorkerConfig": {"agg_workers": 1, "agg_gpus_per_worker": 8}, + "NodeConfig": {"num_gpus_per_node": 8, "system_name": system}, + "SlaConfig": {"isl": 1024, "osl": 256}, + "ModelConfig": {"is_moe": True, "prefix": 0}, + "BenchConfig": {}, + "params": { + "agg": { + "tensor_parallel_size": 8, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + "moe_tensor_parallel_size": 1, + "moe_expert_parallel_size": 8, + "max_batch_size": 32, + "max_num_tokens": 4096, + "max_seq_len": 4096, + }, + }, + } + + +def _moe_backend(version: str, system: str) -> str: + artifacts = generate_backend_artifacts(backend="trtllm", backend_version=version, params=_params(system)) + name = next(name for name in artifacts if "extra_engine_args" in name and "agg" in name) + return yaml.safe_load(artifacts[name])["moe_config"]["backend"] + + +@pytest.mark.parametrize(("version", "expected"), [("1.3.0rc23", "WIDEEP"), ("1.3.0rc29", "AUTO")]) +def test_blackwell_render_keeps_wideep_before_rc25_and_auto_from_rc29(version, expected): + assert _moe_backend(version, "b200_sxm") == expected + + +def test_rc29_hopper_render_keeps_the_hardware_fact(): + assert _moe_backend("1.3.0rc29", "h200_sxm") == "CUTLASS" diff --git a/python/aisimulate/tests/unit/generator/test_trtllm_msa_implementation.py b/python/aisimulate/tests/unit/generator/test_trtllm_msa_implementation.py index bb2ef369c5..2b73242847 100644 --- a/python/aisimulate/tests/unit/generator/test_trtllm_msa_implementation.py +++ b/python/aisimulate/tests/unit/generator/test_trtllm_msa_implementation.py @@ -128,3 +128,60 @@ def test_naive_generator_carries_the_msa_prescription(): mode="agg", ) assert "msa_sparse_implementation" not in off_family["ModelConfig"] + + +# --- rc29: sparse_attention_config.algorithm (required on every SM) --------------------- +_RC29 = "1.3.0rc29" + + +def test_rc29_template_renders_algorithm_for_m3_without_implementation(): + """SM90 M3 (no fmha_sm100 prescription) still needs algorithm=minimax_m3 at rc29.""" + artifacts = generate_backend_artifacts( + backend="trtllm", backend_version=_RC29, params=_params({"msa_sparse_algorithm": "minimax_m3"}) + ) + assert _engine_yaml(artifacts).get("sparse_attention_config") == {"algorithm": "minimax_m3"} + + +def test_rc29_template_renders_both_keys_on_sm100(): + artifacts = generate_backend_artifacts( + backend="trtllm", + backend_version=_RC29, + params=_params({"msa_sparse_algorithm": "minimax_m3", "msa_sparse_implementation": "msa"}), + ) + assert _engine_yaml(artifacts).get("sparse_attention_config") == { + "algorithm": "minimax_m3", + "implementation": "msa", + } + + +def test_rc23_template_ignores_the_algorithm_key(): + """Prior template untouched: rc23 never read `algorithm`, so it must not appear.""" + artifacts = generate_backend_artifacts( + backend="trtllm", backend_version=_BACKEND_VERSION, params=_params({"msa_sparse_algorithm": "minimax_m3"}) + ) + assert "sparse_attention_config" not in _engine_yaml(artifacts) + + +@pytest.mark.parametrize( + ("backend", "model", "expected"), + [ + ("trtllm", "MiniMaxAI/MiniMax-M3", "minimax_m3"), + ("trtllm", "nvidia/MiniMax-M3-NVFP4", "minimax_m3"), + ("trtllm", "deepseek-ai/DeepSeek-V3", None), + ("sglang", "MiniMaxAI/MiniMax-M3", None), + ], +) +def test_algorithm_prescription_is_architecture_keyed_and_trtllm_only(backend, model, expected): + from aisimulate.generator.module_bridge import _msa_sparse_algorithm + + assert _msa_sparse_algorithm(_FakeTask(backend, model, "h200_sxm")) == expected + + +def test_naive_generator_carries_the_algorithm_on_every_system(): + from aisimulate.generator.naive import build_naive_generator_params + + for system in ("b200_sxm", "h200_sxm"): + params = build_naive_generator_params( + model_name="MiniMaxAI/MiniMax-M3", total_gpus=8, system_name=system, backend_name="trtllm", mode="agg" + ) + assert params["ModelConfig"]["msa_sparse_algorithm"] == "minimax_m3" diff --git a/python/aisimulate/tests/unit/generator/test_vllm_kda_decode_batch_cap.py b/python/aisimulate/tests/unit/generator/test_vllm_kda_decode_batch_cap.py new file mode 100644 index 0000000000..40b6992e86 --- /dev/null +++ b/python/aisimulate/tests/unit/generator/test_vllm_kda_decode_batch_cap.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""KDA (Kimi Delta Attention) models on vLLM: the decode batch is capped at 256. + +vLLM 0.29.0 / SM90 faults at the replay of the 512-batch decode CUDA graph +for Kimi-K3 (64/256/384/448/480/496 pass, 512 fails, graphs off at 512 +passes); owner decision 2026-09-25: cap max_batch_size at 256 for models +declaring linear_attn_config.kda_layers. The fact enters the rules as +ModelConfig.has_kda from both render paths (module_bridge, naive). +""" + +import pytest + +from aisimulate.generator.rendering.rule_engine import apply_rule_plugins +from aisimulate.generator.utils import model_has_kda + +pytestmark = pytest.mark.unit + + +def _pv(has_kda, batch=512, preserve=False): + return { + "SlaConfig": {"isl": 4000, "osl": 500}, + "DynConfig": {"mode": "agg"}, + "ModelConfig": {"is_moe": True, "has_kda": has_kda}, + "params": { + "agg": { + "max_batch_size": batch, + "tokens_per_block": 32, + **({"preserve_engine_limits": True} if preserve else {}), + } + }, + } + + +def test_kda_model_decode_batch_capped_at_256_and_budgets_follow(): + pv = _pv(True) + apply_rule_plugins(pv, backend="vllm") + agg = pv["params"]["agg"] + assert agg["max_batch_size"] == 256 + # token budgets derive from the CAPPED batch (256 + 4000 + 1500 = 5756 -> 5760) + assert agg["max_num_tokens"] == 5760 + + +def test_non_kda_model_keeps_the_512_floor(): + pv = _pv(False, batch=256) + apply_rule_plugins(pv, backend="vllm") + assert pv["params"]["agg"]["max_batch_size"] == 512 # the pre-existing floor, untouched + pv = _pv(None, batch=256) # fact absent (older callers): no cap + apply_rule_plugins(pv, backend="vllm") + assert pv["params"]["agg"]["max_batch_size"] == 512 + + +def test_preserve_engine_limits_disables_the_cap(): + pv = _pv(True, batch=512, preserve=True) + apply_rule_plugins(pv, backend="vllm") + assert pv["params"]["agg"]["max_batch_size"] == 512 + + +def test_model_has_kda_reads_the_bundled_config(): + assert model_has_kda("moonshotai/Kimi-K3") is True + assert model_has_kda("Qwen/Qwen3-8B") is False + assert model_has_kda("org/does-not-exist-anywhere") is False diff --git a/python/aisimulate/tests/unit/generator/test_vllm_nvfp4_dsa_kv_cache_dtype.py b/python/aisimulate/tests/unit/generator/test_vllm_nvfp4_dsa_kv_cache_dtype.py new file mode 100644 index 0000000000..9940b328bf --- /dev/null +++ b/python/aisimulate/tests/unit/generator/test_vllm_nvfp4_dsa_kv_cache_dtype.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""NVFP4 DSA checkpoints on vLLM render an explicit fp8 kv-cache dtype. + +vLLM 0.29.0 resolves `--kv-cache-dtype auto` for the modelopt NVFP4 DSA +artifacts (hf_quant kv_cache_quant_algo FP8) to the literal fp8_e4m3 and its +sparse-MLA selector rejects that spelling; the same engine with an explicit +`fp8` loads FlashMLASparseImpl. The prescription is keyed on the checkpoint +architecture and the task's nvfp4 GEMM quant mode — never on model names. +""" + +from __future__ import annotations + +import pytest + +from aisimulate.generator import utils +from aisimulate.generator.rendering.engine import render_backend_parameters + +pytestmark = pytest.mark.unit + +_ARCH = { + "nvidia/GLM-5.3-NVFP4": "GlmMoeDsaForCausalLM", + "nvidia/GLM-5.2-NVFP4": "GlmMoeDsaForCausalLM", + "nvidia/DeepSeek-V3.2-NVFP4": "DeepseekV32ForCausalLM", + "zai-org/GLM-5.3": "GlmMoeDsaForCausalLM", + "nvidia/Llama-3.3-70B-Instruct-FP4": "LlamaForCausalLM", + "local/unbundled": None, +} + + +_MODELOPT_FP8_KV = { + "quant_method": "modelopt", + "quant_algo": "NVFP4", + "kv_cache_scheme": {"dynamic": False, "num_bits": 8, "type": "float"}, +} +# what utils._bundled_quantization returns: config.json quantization block +# merged with the SDK-attached hf_quant_config (kv_cache_quant_algo lives there) +_QUANT = { + "nvidia/GLM-5.2-NVFP4": _MODELOPT_FP8_KV, # config.json carries the KV scheme + "nvidia/GLM-5.3-NVFP4": { + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION", + "kv_cache_quant_algo": "FP8", + }, # hf_quant only + "nvidia/DeepSeek-V3.2-NVFP4": { + "quant_method": "nvfp4", + "quant_algo": "NVFP4", + "kv_cache_quant_algo": "FP8", + }, # hf_quant only + "nvidia/GLM-5-NVFP4-unquant-bundle": {"quant_method": "nvfp4"}, # neither spelling present + "zai-org/GLM-5.3": {"quant_method": "fp8", "fmt": "e4m3", "weight_block_size": [128, 128]}, + "nvidia/Llama-3.3-70B-Instruct-FP4": _MODELOPT_FP8_KV, +} +_ARCH["nvidia/GLM-5-NVFP4-unquant-bundle"] = "GlmMoeDsaForCausalLM" + + +@pytest.fixture(autouse=True) +def _stub_sdk_lookups(monkeypatch): + # the helpers take an optional frozen config (naive path); the stubs ignore it + monkeypatch.setattr(utils, "_model_architecture", lambda model_path, model_config=None: _ARCH.get(model_path)) + monkeypatch.setattr(utils, "_bundled_quantization", lambda model_path, raw=None: _QUANT.get(model_path)) + + +@pytest.mark.parametrize( + ("backend", "model", "gemm", "expected"), + [ + # optimized path: the task's nvfp4 GEMM quant mode is the artifact fact + ("vllm", "nvidia/GLM-5.3-NVFP4", "nvfp4", "fp8"), + ("vllm", "nvidia/DeepSeek-V3.2-NVFP4", "nvfp4", "fp8"), + # naive path (no quant mode): the artifact's fp8 KV fact, either spelling + ("vllm", "nvidia/GLM-5.2-NVFP4", None, "fp8"), # config.json kv_cache_scheme + ("vllm", "nvidia/GLM-5.3-NVFP4", None, "fp8"), # hf_quant kv_cache_quant_algo + ("vllm", "nvidia/GLM-5-NVFP4-unquant-bundle", None, None), # no KV fact anywhere: no prescription + ("vllm", "zai-org/GLM-5.3", "fp8_block", None), # fp8 DSA artifact: auto is fine + ("vllm", "nvidia/Llama-3.3-70B-Instruct-FP4", "nvfp4", None), # NVFP4 but not DSA + ("vllm", "local/unbundled", "nvfp4", None), # unresolvable config: no prescription + ("sglang", "nvidia/GLM-5.2-NVFP4", "nvfp4", None), # vllm-only spelling gap + ], +) +def test_prescription_keyed_on_dsa_architecture_and_an_artifact_fact(backend, model, gemm, expected): + assert utils.vllm_dsa_kv_cache_dtype(backend, model, gemm) == expected + + +def test_artifact_kv_scheme_detection(): + assert utils._artifact_pins_fp8_kv(_MODELOPT_FP8_KV) + assert utils._artifact_pins_fp8_kv({"kv_cache_quant_algo": "FP8"}) + assert not utils._artifact_pins_fp8_kv({"quant_method": "fp8", "fmt": "e4m3"}) + assert not utils._artifact_pins_fp8_kv(None) + + +def test_fp8_kv_dtype_renders_as_the_explicit_vllm_flag(): + out = render_backend_parameters({"kv_cache_dtype": "fp8"}, "vllm") + assert out["kv_cache_dtype"]["kv-cache-dtype"] == "fp8" + # and BF16 renders literally: vLLM >= 0.11 accepts it (backend_config_mapping.yaml guard, #487) + out = render_backend_parameters({"kv_cache_dtype": "bfloat16"}, "vllm") + assert out["kv_cache_dtype"]["kv-cache-dtype"] == "bfloat16" diff --git a/python/aisimulate/tests/unit/generator/test_vllm_token_budget_alignment.py b/python/aisimulate/tests/unit/generator/test_vllm_token_budget_alignment.py new file mode 100644 index 0000000000..a28a4d5899 --- /dev/null +++ b/python/aisimulate/tests/unit/generator/test_vllm_token_budget_alignment.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Every rendered token budget is a multiple of 64, on every backend. + +Engines size per-cache-group slot-mapping rows and scheduler chunks by these +budgets; a non-idiomatic width misaligns every row after the first (vLLM +0.29 DeepSeek-V4: 6012 crashes the cutedsl compress kernel at first forward, +6016 passes). The alignment lives in ONE place — backend_config_mapping.yaml — +so rule formulas stay plain and no backend is a special case (owner decision +2026-09-24). +""" + +import pytest + +from aisimulate.generator.rendering.engine import render_backend_parameters + +pytestmark = pytest.mark.unit + + +def _render(params, backend): + return render_backend_parameters(params, backend) + + +@pytest.mark.parametrize( + ("backend", "param", "key"), + [ + ("vllm", "max_num_tokens", "max-num-batched-tokens"), + ("trtllm", "max_num_tokens", "max_num_tokens"), + ("sglang", "max_prefill_tokens", "max-prefill-tokens"), + ], +) +def test_odd_budget_rounds_up_to_64_on_every_backend(backend, param, key): + assert _render({param: 6012}, backend)[param][key] == 6016 + assert _render({param: 6017}, backend)[param][key] == 6080 # up, never down + assert _render({param: 6016}, backend)[param][key] == 6016 + assert _render({param: 2048}, backend)[param][key] == 2048 + + +def test_sglang_chunked_prefill_size_follows_the_aligned_budget(): + out = _render({"enable_chunked_prefill": True, "max_num_tokens": 6012}, "sglang") + assert out["enable_chunked_prefill"]["chunked-prefill-size"] == 6016 + + +def test_absent_budget_stays_omitted(): + for backend in ("vllm", "trtllm"): + assert "max_num_tokens" not in _render({}, backend) + assert "max_prefill_tokens" not in _render({}, "sglang") + + +@pytest.mark.parametrize("backend", ["vllm", "sglang", "trtllm"]) +def test_rule_engine_aligns_every_budget_after_the_plain_formulas(backend): + """The rule formulas stay plain (vllm agg: 512+4000+1500 = 6012, the DSV4 + crash trigger); the rule engine's single alignment pass rounds every + budget key in every role up to 64 — including trtllm, whose engine-yaml + template bypasses the mapping transforms.""" + from aisimulate.generator.rendering.rule_engine import apply_rule_plugins + + pv = { + "SlaConfig": {"isl": 4000, "osl": 500}, + "DynConfig": {}, + "params": {"agg": {"max_batch_size": 512, "tokens_per_block": 32}}, + } + apply_rule_plugins(pv, backend=backend) + for role, values in pv["params"].items(): + for key in ("max_num_tokens", "max_prefill_tokens"): + if isinstance(values.get(key), int): + assert values[key] % 64 == 0, (backend, role, key, values[key]) + agg = pv["params"]["agg"] + budget_key = "max_prefill_tokens" if backend == "sglang" else "max_num_tokens" + assert agg[budget_key] >= 4000 + 500 # never rounded down (formulas add isl + 500/1500) + + +def test_alignment_pass_applies_without_a_rule_file(): + from aisimulate.generator.rendering.rule_engine import apply_rule_plugins + + pv = {"params": {"agg": {"max_num_tokens": 6012, "max_prefill_tokens": 5500}}} + apply_rule_plugins(pv, backend="no-such-backend") + assert pv["params"]["agg"] == {"max_num_tokens": 6016, "max_prefill_tokens": 5504} diff --git a/python/aisimulate/tests/unit/sdk/database/test_dsv4_sparse.py b/python/aisimulate/tests/unit/sdk/database/test_dsv4_sparse.py index 07041562d6..887d20f8cc 100644 --- a/python/aisimulate/tests/unit/sdk/database/test_dsv4_sparse.py +++ b/python/aisimulate/tests/unit/sdk/database/test_dsv4_sparse.py @@ -5,7 +5,7 @@ Covers: * the per-(attn_kind, mode) module loaders (test-only parsers) - * the sparse-kernel CSV loader (paged_mqa_logits / hca_attn) + * the sparse-kernel CSV loader (paged_mqa_logits) * ``_lookup_sparse_kernel`` (exact + engine resolve + tp fallback) """ @@ -128,14 +128,12 @@ def test_dsv4_test_cases_skipped_under_other_model(monkeypatch): from collector.case_generator import ( get_dsv4_csa_context_test_cases, get_dsv4_csa_generation_test_cases, - get_dsv4_hca_attn_test_cases, get_dsv4_paged_mqa_logits_test_cases, ) assert get_dsv4_csa_context_test_cases() == [] assert get_dsv4_csa_generation_test_cases() == [] assert get_dsv4_paged_mqa_logits_test_cases() == [] - assert get_dsv4_hca_attn_test_cases() == [] @pytest.mark.parametrize( @@ -169,17 +167,11 @@ def test_dsv4_sparse_test_cases_emit_one_kernel_case_per_model(monkeypatch, mode TP is no longer a case axis — the worker fixes tp=1 internally because the kernel is TP-invariant.""" monkeypatch.setenv("COLLECTOR_MODEL_PATH", model_path) - from collector.case_generator import ( - get_dsv4_hca_attn_test_cases, - get_dsv4_paged_mqa_logits_test_cases, - ) + from collector.case_generator import get_dsv4_paged_mqa_logits_test_cases paged = get_dsv4_paged_mqa_logits_test_cases() - hca = get_dsv4_hca_attn_test_cases() assert {c[1] for c in paged} == {"paged_mqa_logits"} - assert {c[1] for c in hca} == {"hca_attn"} assert {c[0] for c in paged} == {model_path} - assert {c[0] for c in hca} == {model_path} # ─────────────────────────────────────────────────────────────────────── diff --git a/python/aisimulate/tests/unit/tools/test_collect_campaign.py b/python/aisimulate/tests/unit/tools/test_collect_campaign.py new file mode 100644 index 0000000000..1b0724b9af --- /dev/null +++ b/python/aisimulate/tests/unit/tools/test_collect_campaign.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""tools/perf_database/collect_campaign.py — the invariants a hand-rolled launcher drifted on (sm120, 2026-09-30).""" + +import importlib.util +import sys +from pathlib import Path + +import pytest +import yaml + +pytestmark = pytest.mark.unit + +_TOOL = Path(__file__).resolve().parents[3] / "tools" / "perf_database" / "collect_campaign.py" + + +def _load(): + spec = importlib.util.spec_from_file_location("collect_campaign", _TOOL) + mod = importlib.util.module_from_spec(spec) + sys.modules["collect_campaign"] = mod + spec.loader.exec_module(mod) + return mod + + +def _plan(tmp_path, shards): + p = tmp_path / "plan.yaml" + p.write_text(yaml.safe_dump({"backend": "vllm", "sm": 120, "shards": shards})) + return p + + +def test_collector_argv_always_resumes_and_never_keeps_csv(tmp_path): + c = _load() + plan = c.load_plan(_plan(tmp_path, [{"op": "gemm", "name": "bf16", "case_filter": "['bfloat16', "}, {"op": "kda"}])) + for shard in plan.shards: + argv = c.collector_argv(plan, shard) + assert "--resume" in argv and "--keep-csv" not in argv + assert argv[argv.index("--sm") + 1] == "120" + assert argv[argv.index("--ops") + 1] == shard.op + assert "--case-filter" in c.collector_argv(plan, plan.shards[0]) + assert "--case-filter" not in c.collector_argv(plan, plan.shards[1]) + assert plan.shards[1].name == "all" + + +def test_one_namespace_per_shard_and_readonly_checkout(tmp_path): + c = _load() + plan = c.load_plan(_plan(tmp_path, [{"op": "gemm", "name": "bf16"}, {"op": "gemm", "name": "fp8"}])) + dirs = {c.shard_dir(tmp_path, s) for s in plan.shards} + assert len(dirs) == 2 + ref = "0123456789abcdef0123456789abcdef01234567" + argv = c.docker_argv( + plan, + plan.shards[0], + image="img", + gpu=3, + checkout=Path("/ck"), + shard_dir=dirs.pop(), + container_name="x", + collector_ref=ref, + ) + assert '"device=3"' in argv and "/ck:/ais:ro" in argv and argv[argv.index("-w") + 1] == "/out" + # the mount has no .git and the images no git: provenance gets the commit from the launcher + assert f"AISIM_COLLECTOR_REF={ref}" in argv + + +def test_collector_ref_is_the_clean_checkout_commit(tmp_path): + import subprocess + + c = _load() + + def git(*args): + subprocess.run(["git", "-C", str(tmp_path), *args], check=True, capture_output=True) + + git("init", "-q") + (tmp_path / "collect.py").write_text("v1\n") + git("add", "collect.py") + git("-c", "user.name=t", "-c", "user.email=t@example.com", "commit", "-q", "-m", "v1") + ref = c.collector_ref(tmp_path) + assert len(ref) == 40 and int(ref, 16) >= 0 + (tmp_path / "collect.py").write_text("v2\n") # uncommitted: the rows would name the wrong code + with pytest.raises(SystemExit, match="uncommitted"): + c.collector_ref(tmp_path) + + +def test_duplicate_shard_keys_are_rejected(tmp_path): + c = _load() + with pytest.raises(ValueError): + c.load_plan(_plan(tmp_path, [{"op": "gemm", "name": "a"}, {"op": "gemm", "name": "a"}])) + + +def test_image_is_the_manifest_pin_plus_finalize_deps(): + c = _load() + base = c.pinned_image("vllm", "gemm") + assert base.startswith("vllm/vllm-openai:") # framework_manifest default pin, never a hand-typed tag + df = c.dockerfile(base) + assert df.startswith(f"FROM {base}\n") and "pyarrow" in df and "pandas" in df + assert c.collect_image_tag(base).startswith("aisim-collect:") + + +def test_shard_state_reads_the_namespace(tmp_path): + c = _load() + plan = c.load_plan(_plan(tmp_path, [{"op": "gemm", "name": "bf16"}])) + s = plan.shards[0] + assert c.shard_state(tmp_path, s) == "pending" + d = c.shard_dir(tmp_path, s) + (d / ".ckpt").mkdir(parents=True) + assert c.shard_state(tmp_path, s) == "started" + (d / "DONE").touch() + assert c.shard_state(tmp_path, s) == "done-no-parquet" # the --keep-csv failure mode, picked up by `finalize` + (d / "gemm_perf.parquet").write_bytes(b"") + assert c.shard_state(tmp_path, s) == "finalized" diff --git a/python/aisimulate/tools/perf_database/collect_campaign.py b/python/aisimulate/tools/perf_database/collect_campaign.py new file mode 100644 index 0000000000..c6a07e4d20 --- /dev/null +++ b/python/aisimulate/tools/perf_database/collect_campaign.py @@ -0,0 +1,449 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run a perf-data collection campaign as declared shards, one container per shard. + +Why this exists (sm120 campaign, 2026-09-30): every box so far re-invented the +same launcher by hand — an image with the finalize deps added, ``--case-filter`` +shards, one cwd per shard so staging CSVs never collide, a GPU pool, and the +finalize step. Each re-invention drifted somewhere (the sm120 run passed +``--keep-csv`` and silently produced no parquet/provenance for 13 shards; the +official image lacks pyarrow so the smoke run died at finalize). This tool fixes +the invariants once: + +* the image is the framework_manifest pin for the (framework, op) being run — + never a hand-typed tag — plus the finalize deps baked on top (``build-image``); +* one output/checkpoint namespace per shard (``///``), exactly + as the collector-upgrade playbook §5 asks; +* ``collect.py`` is always invoked with ``--resume`` and never ``--keep-csv``, + so a finished shard always ends with parquet + collection_meta.yaml; +* every container launch is appended to ``/campaign.jsonl`` (image digest, + argv, exit code, parquet outputs) — the work log the playbook §4 wants. + +Shard plan (YAML):: + + backend: vllm + sm: 120 + shards: + - {op: gemm, name: bf16, case_filter: "['bfloat16', "} + - {op: gemm, name: fp8, case_filter: "['fp8', "} + - {op: kda, name: all} + +Usage:: + + python3 tools/perf_database/collect_campaign.py build-image --backend vllm --op gemm + python3 tools/perf_database/collect_campaign.py run --plan shards.yaml --root /ws/collect --gpus 0,1,2,3 + python3 tools/perf_database/collect_campaign.py status --plan shards.yaml --root /ws/collect + python3 tools/perf_database/collect_campaign.py finalize --plan shards.yaml --root /ws/collect --gpus 0 + +``run`` is re-entrant: shards with a ``DONE`` marker are skipped, unfinished +shards resume from their checkpoint. ``finalize`` re-enters finished shards that +have no parquet (e.g. runs made by hand with ``--keep-csv``) — every task is +already checkpointed, so collect.py only finalizes. Merging the shards of one op +into one table is the delivery step (playbook §8) and is out of scope here. +""" + +from __future__ import annotations + +import argparse +import json +import re +import shlex +import subprocess +import sys +import time +from dataclasses import dataclass +from pathlib import Path + +import yaml + +PKG_ROOT = Path(__file__).resolve().parents[2] # python/aisimulate +if str(PKG_ROOT) not in sys.path: + sys.path.insert(0, str(PKG_ROOT)) + +FINALIZE_DEPS = ("pyarrow", "pandas") # collector/helper.py finalize_perf_files imports +CONTAINER_CHECKOUT = "/ais" +CONTAINER_OUT = "/out" + + +@dataclass(frozen=True) +class Shard: + op: str + name: str + case_filter: str | None = None + + @property + def key(self) -> str: + return f"{self.op}/{self.name}" + + +@dataclass(frozen=True) +class Plan: + backend: str + sm: int + shards: tuple[Shard, ...] + + +def load_plan(path: Path) -> Plan: + doc = yaml.safe_load(path.read_text()) + shards = [] + seen: set[str] = set() + for raw in doc["shards"]: + s = Shard(op=str(raw["op"]), name=str(raw.get("name", "all")), case_filter=raw.get("case_filter") or None) + if s.key in seen: + raise ValueError(f"duplicate shard {s.key}") + seen.add(s.key) + shards.append(s) + return Plan(backend=str(doc["backend"]), sm=int(doc["sm"]), shards=tuple(shards)) + + +def pinned_image(backend: str, op: str) -> str: + """The manifest-pinned image for (backend, op) — family pins included.""" + from collector.framework_manifest import resolve_op_runtime + + return resolve_op_runtime(backend, op).image() + + +def collect_image_tag(base_image: str) -> str: + """Deterministic local tag for + finalize deps.""" + name = base_image.split("@", 1)[0].replace("/", "-").replace(":", "-") + return f"aisim-collect:{name}" + + +def dockerfile(base_image: str) -> str: + return ( + f"FROM {base_image}\n" + "# collector finalize deps the framework image does not ship (collector/helper.py finalize_perf_files)\n" + f"RUN pip install --no-cache-dir {' '.join(FINALIZE_DEPS)}\n" + ) + + +def build_image(base_image: str, *, network: str = "host", extra_build_args: list[str] | None = None) -> str: + tag = collect_image_tag(base_image) + cmd = ["docker", "build", "--network", network, "-t", tag, "-f", "-", "."] + (extra_build_args or []) + subprocess.run(cmd, input=dockerfile(base_image), text=True, check=True) + return tag + + +def image_digest(image: str) -> str | None: + try: + out = subprocess.run( + # a locally built collect image has NO RepoDigests — `index .RepoDigests 0` fails and the launcher + # refused its own build-image output (B200 2026-10-03); fall back to the image Id + ["docker", "inspect", "--format", "{{if .RepoDigests}}{{index .RepoDigests 0}}{{end}}|{{.Id}}", image], + capture_output=True, + text=True, + check=True, + ).stdout.strip() + return out + except Exception: + return None + + +def collector_argv(plan: Plan, shard: Shard) -> list[str]: + """The exact collect.py argv for one shard. ``--resume`` always, ``--keep-csv`` never.""" + argv = [ + f"{CONTAINER_CHECKOUT}/collector/collect.py", + "--backend", + plan.backend, + "--model-cases-full", + "--sm", + str(plan.sm), + "--ops", + shard.op, + "--checkpoint-dir", + f"{CONTAINER_OUT}/.ckpt", + "--resume", + ] + if shard.case_filter: + argv += ["--case-filter", shard.case_filter] + assert "--keep-csv" not in argv + return argv + + +def collector_ref(checkout: Path) -> str: + """The commit the mounted checkout runs, for every shard's provenance. + + The container mounts only python/aisimulate (no .git) and the framework images carry no + git, so collect.py's own lookup would record "unknown"; it reads AISIM_COLLECTOR_REF + instead. A dirty tree is refused: its rows would name a commit that did not produce them. + """ + + def git(*args: str) -> str: + return subprocess.run( + ["git", "-C", str(checkout), *args], capture_output=True, text=True, check=True + ).stdout.strip() + + try: + ref = git("rev-parse", "HEAD") + dirty = git("status", "--porcelain", "--untracked-files=no", "--", ".") + except (OSError, subprocess.CalledProcessError) as error: + raise SystemExit(f"{checkout}: cannot resolve the collector commit ({error})") from error + if dirty: + raise SystemExit(f"{checkout}: uncommitted changes; commit them so the rows record the code that ran") + if not re.fullmatch(r"[0-9a-f]{40,64}", ref): + raise SystemExit(f"{checkout}: `git rev-parse HEAD` gave {ref!r}, not a full commit hash") + return ref + + +def docker_argv( + plan: Plan, + shard: Shard, + *, + image: str, + gpu: int, + checkout: Path, + shard_dir: Path, + container_name: str, + collector_ref: str, +) -> list[str]: + return [ + "docker", + "run", + "--rm", + "--name", + container_name, + "--gpus", + f'"device={gpu}"', + "--shm-size", + "16g", + "-e", + "CUDA_MPS_PIPE_DIRECTORY=/nonexistent-no-mps", + "-e", + "HF_HUB_OFFLINE=1", + "-e", + f"PYTHONPATH={CONTAINER_CHECKOUT}", + "-e", + f"AISIM_COLLECTOR_REF={collector_ref}", + "-e", + f"TRITON_CACHE_DIR={CONTAINER_OUT}/.cache/triton", + "-e", + f"DG_JIT_CACHE_DIR={CONTAINER_OUT}/.cache/deep_gemm", + "-e", + f"FLASHINFER_WORKSPACE_BASE={CONTAINER_OUT}/.cache", + "-v", + f"{checkout}:{CONTAINER_CHECKOUT}:ro", + "-v", + f"{shard_dir}:{CONTAINER_OUT}", + "-w", + CONTAINER_OUT, + "--entrypoint", + "python3", + image, + *collector_argv(plan, shard), + ] + + +def shard_dir(root: Path, shard: Shard) -> Path: + return root / shard.op / shard.name + + +def shard_state(root: Path, shard: Shard) -> str: + d = shard_dir(root, shard) + if not d.exists(): + return "pending" + if list(d.glob("*.parquet")): + return "finalized" + if (d / "DONE").exists(): + return "done-no-parquet" + if (d / ".ckpt").exists(): + return "started" + return "pending" + + +def _log(root: Path, record: dict) -> None: + with (root / "campaign.jsonl").open("a") as fh: + fh.write(json.dumps({"ts": time.strftime("%Y-%m-%dT%H:%M:%S"), **record}) + "\n") + + +def run_shard( + plan: Plan, + shard: Shard, + *, + root: Path, + gpu: int, + checkout: Path, + image: str, + collector_ref: str, + mode: str = "run", +) -> int: + d = shard_dir(root, shard) + (d / ".cache").mkdir(parents=True, exist_ok=True) + name = f"{'f' if mode == 'finalize' else 'c'}_{plan.backend}_{shard.op}_{shard.name}" + cmd = docker_argv( + plan, + shard, + image=image, + gpu=gpu, + checkout=checkout, + shard_dir=d, + container_name=name, + collector_ref=collector_ref, + ) + _log( + root, + { + "event": "start", + "mode": mode, + "shard": shard.key, + "gpu": gpu, + "image": image, + "image_digest": image_digest(image), + "collector_ref": collector_ref, + "argv": collector_argv(plan, shard), + }, + ) + with (d / "run.log").open("a") as log: + log.write(f"### {time.strftime('%FT%T')} {mode} gpu{gpu} {shlex.join(cmd)}\n") + log.flush() + rc = subprocess.run(cmd, stdout=log, stderr=subprocess.STDOUT).returncode + log.write(f"### {time.strftime('%FT%T')} exit={rc}\n") + parquets = sorted(p.name for p in d.glob("*.parquet")) + if rc == 0: + (d / "DONE").touch() + _log(root, {"event": "end", "mode": mode, "shard": shard.key, "gpu": gpu, "exit": rc, "parquet": parquets}) + return rc + + +def _free_gpu(busy: dict[int, subprocess.Popen | None], gpus: list[int]) -> int | None: + for g in gpus: + p = busy.get(g) + if p is None or p.poll() is not None: + return g + return None + + +def run_campaign( + plan: Plan, *, plan_path: Path, root: Path, gpus: list[int], checkout: Path, image_override: str | None, mode: str +) -> int: + """GPU pool over the plan. mode=run: skip DONE shards; mode=finalize: only done-no-parquet shards.""" + root.mkdir(parents=True, exist_ok=True) + todo = [] + for s in plan.shards: + st = shard_state(root, s) + if mode == "run" and st in ("finalized", "done-no-parquet") and (shard_dir(root, s) / "DONE").exists(): + continue + if mode == "finalize" and st != "done-no-parquet": + continue + todo.append(s) + print(f"{mode}: {len(todo)}/{len(plan.shards)} shards on gpus {gpus}") + ref = collector_ref(checkout) # once, before any shard starts: every shard records the same commit + images = {s.op: (image_override or collect_image_tag(pinned_image(plan.backend, s.op))) for s in todo} + for img in sorted(set(images.values())): + if image_digest(img) is None: + raise SystemExit(f"image {img} not present — run `build-image` for its base first") + busy: dict[int, subprocess.Popen | None] = {} + failures = 0 + while todo: + g = _free_gpu(busy, gpus) + if g is None: + time.sleep(15) + continue + finished = busy.get(g) + if finished is not None and finished.returncode != 0: # _free_gpu polled it + failures += 1 + s = todo.pop(0) + print(f"{time.strftime('%T')} gpu{g} <- {s.key}") + busy[g] = subprocess.Popen( + [ + sys.executable, + __file__, + "_shard", + "--plan", + str(plan_path), + "--root", + str(root), + "--gpu", + str(g), + "--shard", + s.key, + "--checkout", + str(checkout), + "--image", + images[s.op], + "--collector-ref", + ref, + "--mode", + mode, + ] + ) + for p in busy.values(): + if p is not None and p.wait() != 0: + failures += 1 + print(f"{mode} finished, {failures} shard(s) exited non-zero") + return 1 if failures else 0 + + +def status(plan: Plan, root: Path) -> None: + for s in plan.shards: + d = shard_dir(root, s) + last = "" + log = d / "run.log" + if log.exists(): + tail = log.read_text(errors="replace").replace("\r", "\n").splitlines() + last = next( + (ln for ln in reversed(tail) if "it/s" in ln or "s/it" in ln or "Completed" in ln or "exit=" in ln), "" + ) + print(f"{s.op:28s} {s.name:20s} {shard_state(root, s):16s} {last[:100]}") + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = ap.add_subparsers(dest="cmd", required=True) + b = sub.add_parser("build-image", help="build + finalize deps") + b.add_argument("--backend", required=True) + b.add_argument("--op", required=True, help="any op of the backend; family pins may select a different image") + b.add_argument("--build-arg", action="append", default=[], help="extra docker build args, e.g. http_proxy=...") + for name in ("run", "finalize", "status"): + p = sub.add_parser(name) + p.add_argument("--plan", required=True, type=Path) + p.add_argument("--root", required=True, type=Path, help="campaign root: one // namespace per shard") + if name != "status": + p.add_argument("--gpus", required=True, help="comma list of GPU indices for the pool") + p.add_argument("--checkout", type=Path, default=PKG_ROOT, help="python/aisimulate checkout to mount") + p.add_argument("--image", default=None, help="override the manifest-derived collect image (A/B only)") + h = sub.add_parser("_shard") # internal: one shard in one container (spawned by run/finalize) + for flag in ("--plan", "--root", "--checkout"): + h.add_argument(flag, required=True, type=Path) + h.add_argument("--gpu", required=True, type=int) + h.add_argument("--shard", required=True) + h.add_argument("--image", required=True) + h.add_argument("--collector-ref", required=True) + h.add_argument("--mode", required=True) + args = ap.parse_args() + if args.cmd == "_shard": + plan = load_plan(args.plan) + shard = next(s for s in plan.shards if s.key == args.shard) + return run_shard( + plan, + shard, + root=args.root, + gpu=args.gpu, + checkout=args.checkout, + image=args.image, + collector_ref=args.collector_ref, + mode=args.mode, + ) + if args.cmd == "build-image": + base = pinned_image(args.backend, args.op) + tag = build_image(base, extra_build_args=[f"--build-arg={a}" for a in args.build_arg]) + print(f"built {tag} from {base}") + return 0 + plan = load_plan(args.plan.resolve()) + if args.cmd == "status": + status(plan, args.root.resolve()) + return 0 + gpus = [int(x) for x in args.gpus.split(",") if x.strip()] + return run_campaign( + plan, + plan_path=args.plan.resolve(), + root=args.root.resolve(), + gpus=gpus, + checkout=args.checkout.resolve(), + image_override=args.image, + mode=args.cmd, + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/audit_pr244_data.py b/scripts/audit_pr244_data.py index d79e54eabd..c588451aa3 100644 --- a/scripts/audit_pr244_data.py +++ b/scripts/audit_pr244_data.py @@ -20,6 +20,11 @@ EVIDENCE = ROOT / "benchmarks/evidence/collector/pr244" SOURCE = "1f29ee459882796db312b8a288c040b83edbb211" SYSTEMS = ("b300_sxm", "gb200", "gb300", "h100_sxm", "h200_sxm") +# Tables PR #244 shipped that were later removed from the tree. The ledger keeps their reported +# hashes (it records what was collected); the inventory check compares the remaining tables. +RETIRED_TABLES = { + "dsv4_hca_attn_module_perf.parquet": "hca_attn FMLA sidecar removed with its SDK reader (#649)", +} def sha256(data): @@ -165,7 +170,8 @@ def build_manifest(): systems = summarize_cases(ledger) for system, summary in systems.items(): tables = {path.name: summarize_table(path) for path in sorted((DATA / system).glob("*/vllm/0.25.0/*.parquet"))} - reported = ledger["systems"][system]["reported_tables"] + reported_tables = ledger["systems"][system]["reported_tables"] + reported = {name: table for name, table in reported_tables.items() if name not in RETIRED_TABLES} require(set(tables) == set(reported), f"Table inventory drift: {system}") for name, table in tables.items(): require(