[pgera]

Port the Q8KV4 (E4M3 query, NVFP4 KV) paged sparse prefill work from vllm-project/MSA dev (PRs 30 and 38)
to nv_dev, adapted to this tree's layout and contract. One commit. Outside the Q8KV4 prefill path only the
QMUL4 toolchain probe moves: into inference/msa_v1/_build_utils.py, shared by the decode attention and
decode indexer JITs.

Changes

Build the prefill for SM107

The extension was gated to SM100 and SM103, and its NVFP4 dequant used the QMUL4 instruction, which SM107 does
not have. Build it for sm_107a as well (CUDA 13.5 or newer, the same rule as the decode), and choose the dequant
path with the toolchain probe the decode already had, now in the shared build utils: QMUL4 where the compiler
accepts it, otherwise a packed-FP16 multiply. The FP16 path gives the same bits: an E2M1 code times an E4M3 scale has
at most five significant bits, so the product is exact in FP16 and the only rounding is the final conversion to
E4M3, exactly as in QMUL4.

K1: remove the stalls around the MMA loop

K1 runs one CTA per (page, head, query split) and streams the queries through it in groups of 128 rows (8 queries
x 16 heads). Profiles on both GPUs showed the MMA waiting on the epilogue and the loader, not on memory. Four changes
address that; together they make K1 10 to 25 % faster on the formal shapes.

• Epilogue. Each thread stores 16 chunks of 16 bytes per group, and for each chunk it recomputed the output row's
  address, its 1/row-sum scale and its validity, then the whole warpgroup synchronized before handing the output
  and statistics stages back to the MMA and softmax. Now compute the row state once per row (a thread holds four
  rows per group) and reuse it for the row's chunks, and let each thread hand its stages back as soon as it holds
  the data in registers, so the MMA refills the stage while the stores drain. One 64-column half of the
  accumulator is processed at a time so the epilogue stays inside its 96-register budget.
• Loader. Each group began with a dependent global load of eight packed query indices followed by a named barrier
  across the loader warps. Now fetch eight groups of indices in one coalesced warp load into a 32-slot ring, let
  every loader warp wait on its own stage, and deepen the Q ring from three to eight stages, so the Q loads run far
  enough ahead of the MMA on long sequences.
• Softmax. Only the query's own page needs the causal mask; history pages are fully visible. Skip the per-element
  select when the warp's rows see all 128 keys (a warp-uniform vote), which removes the select from the common case.
• SM107 MMA. Issue the QK and PV MMAs with the K=64 FP8 instruction, which processes twice the K depth per
  instruction. CUTLASS 4.8 provides it only in the shared-memory/shared-memory form, so the form with A in tensor
  memory, which PV needs, is a local atom (sm107_mma_ts.cuh).

K1 also reads the partial buffer's split stride from the tensor instead of assuming contiguous slabs, and zeroes the
combine's tile counters at its start; both serve the combine below.

Combine: one streaming kernel for every prefill op

The combine merges, for every (query, head) row, the 16 per-page partial outputs and partial log-sum-exps the
forward wrote. It is a pure read of the partial buffer (4.3 GB at 16K tokens), and the previous kernel, which
loaded with per-thread cp.async copies, was limited by instruction issue rather than by memory on both GPUs.
Replace it with a kernel that keeps the memory system busy:

• Treat the rows of the whole batch as one flat list; a CTA owns 64 consecutive rows, which is one contiguous 16 KB
  chunk in each split slab, so a tile of one split is a single bulk-tensor (TMA) copy.
• Stream the splits of a tile through a four-slot shared-memory ring of such copies, issued by one thread and
  tracked by one mbarrier per slot, so loads run ahead of the FP32 accumulation. On SM107, use two slots when the
  partial buffer is interleaved across the locality domains; that was measured faster there.
• The forward's epilogue writes each 32-column group in the order of its STG.128 stores, which is a 4x4 transpose
  of 32-bit pairs. The previous kernel put the tile through shared memory to restore the column order before the
  output store. Now the four lanes that share a group exchange their pairs with two shuffle stages and each lane
  issues one coalesced 128-bit store, without the extra shared-memory pass.
• Load the tile's partial log-sum-exps with one 2-D bulk copy. When the head count is not a multiple of 4 the LSE
  slabs are not 16-byte aligned, which TMA requires; then every thread brings its 4-byte entries in with cp.async.
• Keep BF16/FP16 partials and outputs only, and drop the temperature-LSE and output-scale paths, which have no
  caller in this tree.

The BF16 and Q8KV8 prefill wrappers call the same kernel. The file moves from the training-synced set of
vendor_manifest.json to the inference-local list, since it no longer mirrors the training combine.

Place the partials by memory locality domain

Rubin, B200 and B300 report two memory locality domains through the CUDA driver API, and a read from the other
domain is slower than a local one. With the partial buffer interleaved across the domains, half of every combine
CTA's reads are remote, and the combine does not reach the local read rate.

Add PrefillWorkspace, which owns the split partials of a wrapper (one may also serve several wrappers that run
in order on one stream, sized up front with required_bytes). On a device with more than one domain it builds the
partial buffer from 2 MB pages placed in alternating domains through the CUDA 13.4 virtual memory API and pads every
split slab to a multiple of the stripe, so page p of every slab sits in domain p mod D. The combine deals its
tiles to domains by page; each CTA reads its SM id, looks the domain up in a table the extension probes once per
device, claims the next tile of that domain from a per-domain counter, and steals from the other domain once its own
tiles are gone. K1 zeroes the counters at its start, so no memset node is needed and the launch sequence stays K1
then combine. On one domain, an older driver, or a refused localized allocation the workspace falls back to
ordinary device memory and the kernels run unchanged. MM_SPARSE_PREFILL_Q8KV4_LOCALITY_DOMAINS=0 forces that
layout for comparisons.

Do not port vllm dev's early programmatic-dependent-launch trigger in K1: this tree compiles its griddepcontrol
calls in and triggers at the end of each CTA's MMA work; vllm dev removed the early trigger again after measuring it.

Validation

• Rubin (SM107, CUDA 13.5, CuTe DSL 4.8.0): API contract, direct combine (13 cases), synthetic correctness, CUDA
  Graph replay, real-case smoke, probability contract and shared-source manifest tests, decode attention and decode
  indexer contract tests for the shared probe; full 512-case tier (1024 tests, three-run bitwise repeat).
• B300 (SM103, CUDA 13.5 toolkit, driver 615.71): the same prefill suites at CuTe DSL 4.8.1 (203 passed), with the
  kernel built on the FP16 dequant path (130 passed), and at CuTe DSL 4.5.2, the repository minimum (135 passed);
  full 512-case tier 1024 passed. The QMUL4 and FP16 builds produce bitwise identical outputs on the real smoke shapes.
• Pre-existing failures, identical on nv_dev and untouched by this PR: the decode GQA=8 and decode-indexer
  contract tests on SM107.
• Fingerprints against the base on the real smoke shapes (12 shapes x 2 head counts): B300 bitwise identical; Rubin
  bitwise identical on one part and 23/24 on another, where the single differing element is one BF16 ulp out of
  134 M, caused by the SM107 K=64 MMA accumulation order (the K=32 build reproduces the base bitwise; LSE identical).

Performance

Formal metric: benchmarks/inference/msa_v1/attention/prefill/q8kv4/benchmark.py --suite full (128 production
shapes, CUDA Graph replay of wrapper.run, cold-cache rotation, representative-weighted mean). Base = nv_dev
8b2e229; on Rubin the base also carries the SM107 build enablement (arch gates and the FP16 dequant) so that it runs.

Rubin, HBM 4752 MHz part (SM 2376 MHz, CUDA 13.5 aarch64, CuTe DSL 4.8.0, FP16 dequant)

metric	base	port	ratio
formal weighted mean latency (us)	492.3	389.8	1.263x
aggregate useful TFLOPS	1053	1330	1.263x
per-case speedup geomean / min / max			1.237x / 0.989x / 1.310x
cases slower than 1 %			1

case	batch	total_q	weight	base us	port us	speedup
prefill_c17ccf7fbbb5d1ee5648	1	16384	2477	979.6	767.9	1.276x
prefill_167dfbd500b44961b5cb	1	16384	1347	991.3	779.9	1.271x
prefill_34d5c963ddbc73bd8549	1	16384	1010	965.6	767.5	1.258x
prefill_7e6aa8ec3387a0ecc03b	4	9421	198	609.8	480.9	1.268x
prefill_ca3cf58ae0ea6463c340	8	7464	192	487.8	383.8	1.271x

B300 SXM6 (SM103, CUDA 13.5 toolkit, driver 615.71, CuTe DSL 4.8.1, QMUL4 dequant)

metric	base	port	ratio
formal weighted mean latency (us)	931.2	786.2	1.184x
aggregate useful TFLOPS	557	659	1.184x
per-case speedup geomean / min / max			1.155x / 0.985x / 1.231x
cases slower than 1 %			1

case	batch	total_q	weight	base us	port us	speedup
prefill_c17ccf7fbbb5d1ee5648	1	16384	2477	1880.5	1573.6	1.195x
prefill_167dfbd500b44961b5cb	1	16384	1347	1908.5	1582.1	1.206x
prefill_34d5c963ddbc73bd8549	1	16384	1010	1882.7	1574.4	1.196x
prefill_7e6aa8ec3387a0ecc03b	4	9421	198	1109.9	954.2	1.163x
prefill_ca3cf58ae0ea6463c340	8	7464	192	875.9	762.0	1.149x

The one case below 1.0x on each GPU is a shape under 25 us, inside the suite's run-to-run noise (two port runs
differ by 1.000x weighted, 0.978x to 1.006x per case).

🤖 Generated with Claude Code