Repository navigation
Conversation
Indexer: run(..., pdl=True) launches the score kernel with programmatic stream serialization. Before griddepcontrol.wait the kernel reads only the plan, the page table and history K pages (pages strictly before the page that holds the request's first query token) and issues their TMA. Q and the pages that may hold this step's tokens load after the wait. Callers must not set pdl when the immediately preceding kernel wrote the plan or the page table. Decode attention: MSA_Q8KV4_DECODE_ATTN_PDL=1 adds the PDL launch attribute. The kernel waits after its setup (barrier init, TMEM allocation, tensor-map prefetch) and triggers dependents right after setup, so a PDL-launched consumer of the output can launch and prefetch its weights early. Default behavior is unchanged: without a PDL launch griddepcontrol.wait returns immediately. Outputs are bit-identical with the switches on and off. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Orien-Xu
marked this pull request as ready for review
October 24, 2026 00:47
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Hi, I'm Orien from Parasail. This fixes an issue we ran into while serving MiniMax-M3.1-Flash (#25).
Closes #25.
Change
Indexer (
indexer/decode/q8kv4)run(..., pdl: bool = False), plumbed through the binding,IndexerGemmArguments.pdlandIndexerGemmRunner::run. Withpdl=Truethe kernel launches withcudaLaunchAttributeProgrammaticStreamSerialization.sm100_indexer_gemm_umma.hppthe transport warp issues the first-page TMA for history pages(
logical < (kv_length - query_length) / kPageTokens) beforegriddepcontrol.wait; the otherpages and Q load after it. All warps then wait, and the CTA calls
griddepcontrol.launch_dependents.Decode attention (
attention/decode/q8kv4)MSA_Q8KV4_DECODE_ATTN_PDL=1(read once) adds the PDL launch attribute and setsparams.pdl_early_trigger.run_fwd_mainloop_devicecallsgriddepcontrol.waitafter barrier init, TMEM allocation andtensor-map prefetch, and with the early trigger calls
cudaTriggerProgrammaticLaunchCompletion()right after setup.
Without a PDL launch
griddepcontrol.waitreturns immediately, so the default path differs only inthe order of the indexer's first TMA issues. Outputs are bit-identical.
Results
B200, 4-way TP serving, 8-token speculative verify, engine A/B (same GPUs, ABAB, 4 runs per arm):
chain wall per sparse layer -2.8 us at 40 verify rows (5 x 120K) and -3.6 us at 96 rows
(12 x 60K), about 1.3-1.4% of the decode step. Golden outputs identical on/off.
Tests
tests/inference/msa_v1/indexer/decode/q8kv4/test_pdl.py:pdl=Truevspdl=False,eager and CUDA Graph, mixed lengths, bit-equal.
tests/inference/msa_v1/indexer/decode/q8kv4andtests/inference/msa_v1/attention/decode/q8kv4on SM100/SM103 on this base (the patch wasrebased from bc5f83a onto 35dc1be; the only conflict was context in the launch code), plus the
attention suite with
MSA_Q8KV4_DECODE_ATTN_PDL=1, and re-measure on currentnv_dev.GPU validation (NVIDIA B200, driver 595.71.05, CUDA 13.3, torch 2.14.1+cu130, cutlass-dsl 4.8.1, head 175038d, base nv_dev 35dc1be)
Correctness:
indexer/decode/q8kv4/test_pdl.py(pdl on vs off, eager + CUDA Graph, B=1/32/129): 6 passed, bit-identical.MINIMAX_INFERENCE_TEST_SUITE=smokeindexer decode q8kv4: 78 passed. The 1 failure wastest_api_contractpinningrun()'s signature; 175038d updates it, and the rerun gives 17/17 passed fortest_api_contract+test_pdl.MSA_Q8KV4_DECODE_ATTN_PDL=0and with=1.Speed (
pdl_bench.py). Per sparse layer: Hq=16, Hkv=1, index heads=1, 8 query tokens per request, top-16, page 128. Every layer has its own cold KV pool, 8 layers per CUDA Graph, median of 11 x 20 replays. 4 runs per arm in ABBA order. Wall us/layer, off -> on, with output digests (TopK + attention + projection) identical on/off for every case:Chain B, where the neighbours are the PDL-launched MSA kernels, saves 0.8-1.8 us/layer at every batch up to 64, with no regression at large batch. In chain A the neighbours (a torch copy and a cuBLAS GEMM) trigger no early launch, so the effect is within +-0.8% (max +0.14 us).
torch-profiler per-kernel device time, us/layer (
gpu/t3p/prof.jsonl, traces ingpu/t3p/traces/). With PDL, a kernel's interval starts at its early launch and includes waiting atgriddepcontrol.wait, so per-kernel times grow while the wall shrinks: