BUG: qwen3.8-27b export fails warming compiled forwards on SM120

Inizia da ignition/components/paged_tma.py, in particolare da State.run, e confronta i suoi argomenti del kernel con la firma FFI SM120 descritta nell’issue. L’issue segnala che sia le chiamate gated sia quelle ungated omettono l’argomento finale kv_tile_ranges e propone di passare None. Verifica la correzione riscaldando i forwards compilati per l’export qwen3.8-27b segnalato nell’issue su SM120; il lavoro è completato quando tutti gli 81 forwards sono stati riscaldati correttamente.

[ignition] qwen3.8-27b export fails warming compiled forwards on SM120: paged_tma misses trailing kv_tile_ranges argument (Expects 16 parameters when calling _launch)

Summary

Thanks for the lovely work. I have gotten qwen3.8-flash working successfully, but qwen3.8-27b has failed on the prepare step. I did some discovery to identify the issue and have uncovered the bug and created a suggested fix.

ignition export qwen3.8-27b --gpus 0 --max-concurrency 4 on SM120 fails at "Warm compiled forwards" (0/81, ~5 s):

[  FAILED  ] Warm compiled forwards.  0/81  5.3s
error: preparation failed: RuntimeError: Inductor preparation failed across ranks: ['TypeError: Expects 16 parameters when calling: `_launch(q: Tensor([n0, 24, 256], float8_e4m3fn), k: Tensor([n1, 4, 128, 256], float8_e4m3fn), v: Tensor([n1, 4, 128, 256], float8_e4m3fn), o: Tensor([n0, 24, 256], bfloat16), lse: None, softmax_scale_log2: float32, output_scale: float32, skip_softmax_threshold: None, seqlens_kv: Tensor([n2], int32), cu_seqlens_q: Tensor([n3], int32), block_tables: Tensor([n4, n5], int32), cu_seqlens_k: None, max_seqlen_q: int32, use_pdl: bool, output_gate: Tensor([n0, 24, 256], bfloat16), kv_tile_ranges: None)`']

The failing kernel is the gated SM120 FP8 paged-TMA prefill kernel (per the FFI signature: GQA 24:4, head_dim 256, paged 128 tokens/page, FP8 Q/K/V with BF16 output, concrete output_gate), at the first forward in warm order (0/81).

Root-cause analysis

The kernel's compiled FFI entry (CuTe-DSL tvm_ffi_builder arg-count check) strictly expects 16 positional arguments:

q, k, v, o, lse, softmax_scale_log2, output_scale, skip_softmax_threshold,
seqlens_kv, cu_seqlens_q, block_tables, cu_seqlens_k, max_seqlen_q, use_pdl,
output_gate, kv_tile_ranges

stream is not among them (the entry is compiled with the stream bound from the environment).

The failing recipe is qwen38_dense (ignition/recipes/qwen38_dense/), whose prefill attention goes through ignition/components/paged_tma.py (the PreparedTma TMA-prefill path, compiled via flashinfer.cute_dsl.attention.fmha.sm120.compile_sm120_fmha_fp8_paged_kernel with with_output_gate=geometry.gated and no with_kv_tile_ranges). The call site State.run (lines 44-47) builds 14 base arguments (ending at use_pdl=False) and appends output_gate only when the query is gated — 15 for the gated case. The trailing kv_tile_ranges argument (None for a non-split-KV launch) is never passed, so the FFI arg-count check rejects the call.

The ungated path is affected the same way, and this is now confirmed rather than inferred: the fork's own internal caller (flashinfer/attention/cute_dsl/sm120_fmha.py:370-385) invokes the ungated entry with 15 arguments ending in None # kv_tile_ranges — so the ungated FFI entry also expects a trailing kv_tile_ranges (15 total), while hyperflare's ungated call passes 14. Both call-site branches are therefore short by exactly one argument.

Failure path (traceback)

The TypeError is raised on the rank while the warm phase runs the compiled forward; the warm loop captures it as a string and re-raises it cross-rank:

qwen38_dense/bootstrap.py:349   capture_forwards(forwards, signatures, ...)
  → capture.py:74               cache.prepare_many(...)
    → aot.py:807                prepare_many(...)
      → aot.py:869              _prepare_many → self._warm(bound, ...)
        → aot.py:1192           _warm → selections.warm(bound, warmup, ...)
          → inductor_tuning.py:253  warm → self._check(failure)   # failure = "TypeError: Expects 16 parameters …"
            → inductor_tuning.py:133  _check → raise RuntimeError(f"Inductor preparation failed across ranks: {failures}")

So the compile phase succeeds (81/81) and the failure is purely at the first warm run of the gated prefill forward — a kernel-interface arg-count mismatch, not a compile or capture error.

Suggested fix

State.run is missing the trailing kv_tile_ranges argument. Append None (a non-split-KV launch does not use it) in both branches:

-        self.kernel(*arguments, *((output_gate,) if output_gate is not None else ()))
+        self.kernel(*arguments, *((output_gate, None) if output_gate is not None else (None,)))

• Gated: 14 base args + output_gate + None = 16, matching the FFI entry.
• Ungated: 14 base args + None = 15, matching the ungated entry — the fork's own caller (flashinfer/attention/cute_dsl/sm120_fmha.py:370-385) passes 15 args ending in None # kv_tile_ranges.

Validated end-to-end on SM120 (GB202): after this change, ignition export qwen3.8-27b completes "Warm compiled forwards" (81/81) and the model serves — OpenAI /v1/models returns the model. No other changes.

Environment

• hyperflare-inference: 7f534bb (main)
• Locked preparation environment: torch 2.14.1+cu132, triton 3.8.0, flashinfer-python 0.7.3 (fork local-inference-lab/flashinfer @ 85ea6e7, b12x extra), transformers 5.21.0.dev0
• GPU: GB202 (SM120), CUDA 13.2.1, driver 595.104.02
• Model: qwen3.8-27b (checkpoint revision 482ca0f3)

Reproduction

>>> ignition export qwen3.8-27b --gpus 0 --max-concurrency 4
[    OK    ] Resolve checkpoint.  qwen3.8-27b · 482ca0f3 · cached  0.0s
[    OK    ] Select GPUs.  GPU 0 · concurrency 4 · images 0 · up to 7 draft tokens · speculative token capacity 128 · packed graph budget 64 · prompt logprobs off · mixed prefill off · expert scales native · compile workers…  0.3s
Fetching 2 files: 100%|███████████████████████████| 2/2 [00:00<00:00, 833.20it/s]
Download complete:  |  0.00B
Reconstruction complete: |                          |  0.00B /  0.00B
[    OK    ] Inspect checkpoint.  0.0s
[    OK    ] Load weights.  23.8s
[    OK    ] Load draft weights.  3.9s
[    OK    ] Allocate model state.  0.1s
[    OK    ] Construct graph catalog.  13.8s
[    OK    ] Prepare and tune kernels.  44.0s
Graph catalog: 99 graphs (before compilation/capture)
 Kind         Spec  B      W      P                                      D  Prefill  Count
 prefill      on    1      128,256,512,1024,1536,2048,2560,3072,3584,4096  -                                      -  -        10
 prefill      off   1      128,256,512,1024,1536,2048,2560,3072,3584,4096  -                                      -  -        10
 decode       off   1,2,4  1      -                                      -  -        3
 speculative  on    4      8      5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32  7  -        28
 speculative  on    2      8      3,4,5,6,7,8,9,10,11,12,13,14,15,16     7  -        14
 speculative  on    1      8      2,3,4,5,6,7,8                          7  -        7
 speculative  on    4      1      4                                      -  -        1
 speculative  on    2      1      2                                      -  -        1
 speculative  on    1      1      1                                      -  -        1
 resume       on    1,2,4  1      -                                      -  -        3
 init         on    1,2,4  1      -                                      -  -        3
 restore      on    1,2,4  1      -                                      -  -        3
 resume       off   1,2,4  1      -                                      -  -        3
 init         off   1,2,4  1      -                                      -  -        3
 restore      off   1,2,4  1      -                                      -  -        3
 snapshot     off   1,2,4  1      -                                      -  -        3
B=batch  W=logical width/tokens  P=packed rows  D=draft depth
[    OK    ] Compile model forwards.  81/81  161.0s
[  FAILED  ] Warm compiled forwards.  0/81  5.3s
error: preparation failed: RuntimeError: Inductor preparation failed across ranks: ['TypeError: Expects 16 parameters when calling: `_launch(...)`']

Relation to prior issues

• #183 ("compiled forward has unbound static Inductor kernels", fixed by #185): a different failure — a bundle-matching ValueError at forward 26/108 on a 108-graph catalog. This is a kernel-interface TypeError at forward 0/81 on a 99-graph catalog (--max-concurrency 4).
• #187 (warmup OOM on 96 GB SM120): a different failure mode (memory accumulation), later in warmup.

Notes

• The warm phase does not print the failing forward's name (the warm loop suppresses per-forward progress), so the kernel is identified from its FFI signature rather than a graph name; it is a prefill-graph kernel, first in warm order.
• The pinned flashinfer fork (local-inference-lab/flashinfer @ 85ea6e7) defines the kernel interface; the call-site gap is in ignition/components/paged_tma.py.