@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 94e26dd3dd

[Qwen3.8-Flash-Next] Support FP8 indexer cache for QSA (#54890)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

Thien Tran committed Sep 7, 2026 at 20:15 UTC 94e26dd3dd7d6363b524c96521c78feb501bdc61

5 files changed +142 -24

tests/models/qwen4_exp/test_qsa_pre_indexer.py

+31 -6

@@ -49,8 +49,21 @@ def _make_block_table(block_counts):

49 return table, num_blocks

52 +def assert_fp8_within_one_ulp(actual: torch.Tensor, expected: torch.Tensor) -> None:

53 + # e4m3 is sign-magnitude, so within a sign the uint8 code order matches the

54 + # value order and one ulp is one code step. The two paths' intermediates

55 + # differ in the pooling accumulation order, which at denormal magnitudes

56 + # (absolute grid step 2^-9) shows up as up to 2 code steps.

57 + code_diff = (

58 + actual.view(torch.uint8).int() - expected.view(torch.uint8).int()

59 + ).abs()

60 + abs_diff = (actual.float() - expected.float()).abs()

61 + assert bool(((code_diff <= 1) | (abs_diff <= 2**-8)).all())

62 +

63 +

64 @requires_qsa_kernels

65 @pytest.mark.usefixtures("default_vllm_config")

66 +@pytest.mark.parametrize("indexer_dtype", [torch.bfloat16, torch.float8_e4m3fn])

67 @pytest.mark.parametrize(

68 "mrope,is_2d_positions,cache_rope_positions,state_size,seq_lens,query_lens,history_lens",

69 [

@@ -74,6 +87,7 @@ def _make_block_table(block_counts):

87 ],

88 )

89 def test_qsa_fused_pre_indexer_matches_unfused(

90 + indexer_dtype,

91 mrope,

92 is_2d_positions,

93 cache_rope_positions,

@@ -215,7 +229,7 @@ def test_qsa_fused_pre_indexer_matches_unfused(

229 fused_compressed_storage = torch.zeros(

230 num_compressed_blocks,

231 compressed_page_elements + 16,

218 - dtype=torch.bfloat16,

232 + dtype=indexer_dtype,

233 device=device,

234 )

235 fused_compressed = torch.as_strided(

@@ -231,7 +245,7 @@ def test_qsa_fused_pre_indexer_matches_unfused(

245 q_weight = torch.randn(D, dtype=torch.bfloat16, device=device) * 0.2

246 k_weight = torch.randn(D, dtype=torch.bfloat16, device=device) * 0.2

247

234 - fused_query = torch.empty(num_tokens, HQ, D, dtype=torch.bfloat16, device=device)

248 + fused_query = torch.empty(num_tokens, HQ, D, dtype=indexer_dtype, device=device)

249 qsa_pre_indexer(

250 projected_qk[:, : HQ * D],

251 projected_qk[:, HQ * D :],

@@ -290,8 +304,19 @@ def test_qsa_fused_pre_indexer_matches_unfused(

304 if rope_positions is not None:

305 qsa_store_cache_rows(rope_positions, raw_slots, position_rows)

306

293 - torch.testing.assert_close(fused_query, unfused_query, rtol=RTOL, atol=ATOL)

307 + if indexer_dtype == torch.float8_e4m3fn:

308 + # Both paths round the same ~bf16 intermediates to e4m3. Bitwise

309 + # equality (the dsv4 indexer test's bar) does not hold here: the 1D

310 + # reference rope (forward_cuda) differs from _norm_rope by up to 1

311 + # bf16 ulp, and even the MRoPE pair flips a code occasionally.

312 + unfused_query = unfused_query.to(indexer_dtype)

313 + assert_fp8_within_one_ulp(fused_query, unfused_query)

314 + else:

315 + torch.testing.assert_close(fused_query, unfused_query, rtol=RTOL, atol=ATOL)

316 assert torch.equal(fused_raw.view(torch.int16), unfused_raw.view(torch.int16))

295 - torch.testing.assert_close(

296 - fused_compressed, unfused_compressed, rtol=RTOL, atol=ATOL

297 - )

317 + if indexer_dtype == torch.float8_e4m3fn:

318 + assert_fp8_within_one_ulp(fused_compressed, unfused_compressed)

319 + else:

320 + torch.testing.assert_close(

321 + fused_compressed, unfused_compressed, rtol=RTOL, atol=ATOL

322 + )

tests/models/qwen4_exp/test_qsa_reference.py

+61 -6

@@ -48,6 +48,7 @@ def test_qsa_mtp_index_share_updates_cache_but_skips_selection(

48 index_n_heads=1,

49 index_kv_heads=1,

50 index_head_dim=1,

51 + indexer_dtype=torch.bfloat16,

52 raw_key_cache=SimpleNamespace(

53 kv_cache=torch.empty(0),

54 rope_position_cache=None,

@@ -484,6 +485,7 @@ def test_qsa_unfused_cache_update_ignores_padded_qk() -> None:

485 use_fused_pre_indexer=False,

486 index_n_heads=1,

487 index_head_dim=64,

488 + indexer_dtype=torch.bfloat16,

489 q_layernorm=norm,

490 k_layernorm=norm,

491 rotary_emb=rope,

@@ -666,13 +668,16 @@ def test_qsa_fused_metadata_matches_pytorch_for_large_padded_prefill() -> None:

668 (4, 33),

669 ],

670 )

671 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn])

672 def test_qsa_decode_selection_correctness(

670 - decode_query_len: int, num_requests: int

673 + decode_query_len: int, num_requests: int, dtype: torch.dtype

674 ) -> None:

675 torch.manual_seed(1)

676 heads, head_dim = 4, 128

677 rows = num_requests * decode_query_len

675 - q = torch.randn(rows, heads, head_dim, device="cuda", dtype=torch.bfloat16)

678 + q = torch.randn(rows, heads, head_dim, device="cuda", dtype=torch.bfloat16).to(

679 + dtype

680 + )

681 page_size, pages_per_request, max_sequence_length = (

682 (16, 40, 2560) if num_requests > 32 else (4, 20, 320)

683 )

@@ -684,7 +689,7 @@ def test_qsa_decode_selection_correctness(

689 head_dim,

690 device="cuda",

691 dtype=torch.bfloat16,

687 - )

692 + ).to(dtype)

693 page_table = torch.randperm(num_pages, device="cuda", dtype=torch.int32).reshape(

694 num_requests, pages_per_request

695 )

@@ -736,14 +741,39 @@ def test_qsa_decode_selection_correctness(

741 compress_ratio,

742 )

743

744 + if dtype == torch.float8_e4m3fn:

745 + # fp8 logits tie at the top-k boundary more often than bf16, so index

746 + # identity is not stable; compare the selected value multisets.

747 + # SM90 wgmma accumulates fp8 in reduced precision (~3e-4 abs

748 + # observed); SM100 tcgen05 is exact fp32.

749 + rtol = atol = 1e-3 if current_platform.is_device_capability(90) else None

750 + logits = _qsa_mqa_paged_reference(

751 + q, cache, page_table, token_to_req, visible_blocks

752 + )

753 + for row in range(rows):

754 + selected = actual[row][actual[row] >= 0]

755 + wanted = expected[row][expected[row] >= 0]

756 + assert selected.numel() == wanted.numel()

757 + torch.testing.assert_close(

758 + logits[row, selected.long()].sort().values,

759 + logits[row, wanted.long()].sort().values,

760 + rtol=rtol,

761 + atol=atol,

762 + )

763 + return

764 +

765 torch.testing.assert_close(actual.sort().values, expected.sort().values)

766

767

768 @requires_qsa_kernels

769 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn])

770 @pytest.mark.parametrize("seq_len_slack", [0, 1792])

771 @pytest.mark.parametrize("force_chunk", [False, True])

772 def test_qsa_prefill_selection_correctness(

746 - monkeypatch: pytest.MonkeyPatch, seq_len_slack: int, force_chunk: bool

773 + monkeypatch: pytest.MonkeyPatch,

774 + seq_len_slack: int,

775 + force_chunk: bool,

776 + dtype: torch.dtype,

777 ) -> None:

778 # page_size=24 (does not divide the 64-aligned clipped width) and an

779 # oversized page table, so the clipped logits width comes from

@@ -755,8 +785,12 @@ def test_qsa_prefill_selection_correctness(

785 torch.manual_seed(2)

786 query_lens = [3, 33]

787 rows, heads, head_dim = sum(query_lens), 4, 128

758 - q = torch.randn(rows, heads, head_dim, device="cuda", dtype=torch.bfloat16)

759 - cache = torch.randn(128, 24, 1, head_dim, device="cuda", dtype=torch.bfloat16)

788 + q = torch.randn(rows, heads, head_dim, device="cuda", dtype=torch.bfloat16).to(

789 + dtype

790 + )

791 + cache = torch.randn(128, 24, 1, head_dim, device="cuda", dtype=torch.bfloat16).to(

792 + dtype

793 + )

794 page_table = torch.randperm(128, device="cuda", dtype=torch.int32).reshape(2, 64)

795 token_to_req = torch.repeat_interleave(

796 torch.arange(2, device="cuda", dtype=torch.int32),

@@ -804,6 +838,27 @@ def test_qsa_prefill_selection_correctness(

838 compress_ratio,

839 )

840

841 + if dtype == torch.float8_e4m3fn:

842 + # fp8 logits tie at the top-k boundary more often than bf16, so index

843 + # identity is not stable; compare the selected value multisets.

844 + # SM90 wgmma accumulates fp8 in reduced precision (~3e-4 abs

845 + # observed); SM100 tcgen05 is exact fp32.

846 + rtol = atol = 1e-3 if current_platform.is_device_capability(90) else None

847 + logits = _qsa_mqa_paged_reference(

848 + q, cache, page_table, token_to_req, visible_blocks

849 + )

850 + for row in range(rows):

851 + selected = actual[row][actual[row] >= 0]

852 + wanted = expected[row][expected[row] >= 0]

853 + assert selected.numel() == wanted.numel()

854 + torch.testing.assert_close(

855 + logits[row, selected.long()].sort().values,

856 + logits[row, wanted.long()].sort().values,

857 + rtol=rtol,

858 + atol=atol,

859 + )

860 + return

861 +

862 torch.testing.assert_close(actual.sort().values, expected.sort().values)

863

864

vllm/models/qwen4_exp/common/qsa_cache.py

+15 -5

@@ -725,10 +725,16 @@ class QSAMetadataBuilder(AttentionMetadataBuilder[QSAForwardMetadata]):

725

726

727 class QSAStateBackend(AttentionBackend):

728 - """Key-only dummy backend for out-of-band BF16 QSA side-cache operations."""

728 + """Key-only dummy backend for out-of-band QSA side-cache operations."""

729

730 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]

731 - supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16"]

731 + # fp8 entries allow the optional e4m3 compressed indexer cache.

732 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

733 + "auto",

734 + "bfloat16",

735 + "fp8",

736 + "fp8_e4m3",

737 + ]

738

739 @staticmethod

740 def get_name() -> str:

@@ -790,8 +796,12 @@ class _QSAStateCache(nn.Module, AttentionLayerBase):

796 """Adapt the unified [B, H, N, C] view to QSA's [B, N, H, C]."""

797 if kv_cache.ndim != 4 or kv_cache.shape[1] != 1:

798 raise ValueError("QSA state cache must be [blocks, 1, states, width]")

793 - if kv_cache.dtype != torch.bfloat16 or kv_cache.shape[3] != self.head_size:

794 - raise ValueError("QSA state cache does not match its packed BF16 spec")

799 + if kv_cache.dtype != self.dtype or kv_cache.shape[3] != self.head_size:

800 + raise ValueError(

801 + f"QSA state cache does not match its spec "

802 + f"(dtype {kv_cache.dtype} != {self.dtype} or width "

803 + f"{kv_cache.shape[3]} != {self.head_size})"

804 + )

805 super().bind_kv_cache(kv_cache.transpose(1, 2))

806

807 def get_attn_backend(self) -> type[AttentionBackend]:

@@ -850,7 +860,7 @@ class QSAKeyStateCache(_QSAStateCache):

860

861

862 class QSACompressedKeyCache(_QSAStateCache):

853 - """Normalized, group-first-RoPE BF16 key at one row per complete group."""

863 + """Normed, group-first-RoPE key at one row per complete group."""

864

865 def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:

866 del vllm_config

vllm/models/qwen4_exp/nvidia/indexer_qsa.py

+18 -1

@@ -147,6 +147,21 @@ class QSAIndexer(nn.Module):

147

148 cache_config = vllm_config.cache_config

149 cache_prefix = f"{prefix}." if prefix else ""

150 + # Plain e4m3 without scales: Q and the compressed K are RMSNormed

151 + # before quantization, and the logits kernels dot fp8 x fp8 directly.

152 + self.indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype(

153 + "bf16"

154 + )

155 + if self.indexer_kv_dtype == "fp8":

156 + indexer_dtype = torch.float8_e4m3fn

157 + elif self.indexer_kv_dtype == "bf16":

158 + indexer_dtype = torch.bfloat16

159 + else:

160 + raise NotImplementedError(

161 + f"indexer_kv_dtype={self.indexer_kv_dtype!r} is not supported "

162 + "by the Qwen4Exp QSA indexer (only 'bf16' or 'fp8')."

163 + )

164 + self.indexer_dtype = indexer_dtype

165 self.raw_key_cache = QSAKeyStateCache(

166 head_size=self.index_head_dim,

167 dtype=torch.bfloat16,

@@ -158,7 +173,7 @@ class QSAIndexer(nn.Module):

173 )

174 self.compressed_key_cache = QSACompressedKeyCache(

175 head_size=self.index_head_dim,

161 - dtype=torch.bfloat16,

176 + dtype=indexer_dtype,

177 compress_ratio=self.compress_ratio,

178 prefix=f"{cache_prefix}compressed_key_cache",

179 cache_config=cache_config,

@@ -276,6 +291,7 @@ class QSAIndexer(nn.Module):

291 num_tokens,

292 self.index_n_heads,

293 self.index_head_dim,

294 + dtype=self.indexer_dtype,

295 )

296 qsa_pre_indexer(

297 projected_q,

@@ -313,6 +329,7 @@ class QSAIndexer(nn.Module):

329 self.q_layernorm.variance_epsilon,

330 ).reshape_as(q)

331 q = apply_qsa_rope(self.rotary_emb, positions, q)

332 + q = q.to(self.indexer_dtype)

333

334 raw_key_cache = raw_key_state_cache.key_cache

335 rope_position_cache = raw_key_state_cache.rope_position_cache

vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py

+17 -6

@@ -200,6 +200,9 @@ def _qsa_mqa_paged_prefill_kernel(

200 other=0.0,

201 eviction_policy="evict_first",

202 )

203 + # With an fp8 cache, SM90 lowers this dot to reduced-precision wgmma

204 + # accumulation; SM100 tcgen05 accumulates in fp32. The SM90 error

205 + # (~1e-4 relative) is far below the e4m3 quantization noise floor.

206 scores = tl.dot(keys, query, out_dtype=tl.float32)

207 scores = tl.reshape(scores, (BLOCK_N, TILE_R, NUM_HEADS_PADDED))

208 score = tl.sum(tl.maximum(scores, 0.0), axis=2)

@@ -337,7 +340,7 @@ def warmup_qsa_mqa_paged_decode(

340 for decode_query_len, num_requests in profiles:

341 num_rows = decode_query_len * num_requests

342 q_ptr = TritonWarmupTensor(

340 - torch.bfloat16,

343 + k_cache.dtype,

344 shape=(num_rows, num_heads, head_dim),

345 )

346 logits_ptr = TritonWarmupTensor(

@@ -365,7 +368,8 @@ def warmup_qsa_mqa_paged_decode(

368 BLOCK_N=_DECODE_BLOCK_N,

369 TILES_PER_PROG=tiles_per_program,

370 STAGES=2,

368 - num_warps=2,

371 + # tuned on GB300

372 + num_warps=1 if k_cache.dtype == torch.float8_e4m3fn else 2,

373 grid=(

374 num_requests,

375 triton.cdiv(columns, _DECODE_BLOCK_N * tiles_per_program),

@@ -393,7 +397,11 @@ def _prefill_logits(

397 logits = torch.empty(

398 (num_queries, logits_width), dtype=torch.float32, device=q.device

399 )

396 - TILE_R = 64

400 + # tuned on GB300

401 + if k_cache.dtype == torch.float8_e4m3fn:

402 + TILE_R, STAGES, num_warps = 32, 2, 8

403 + else:

404 + TILE_R, STAGES, num_warps = 64, 2, 4

405 BLOCK_N = 64

406 K_TILES = 16

407 grid = (

@@ -421,8 +429,8 @@ def _prefill_logits(

429 TILE_R=TILE_R,

430 BLOCK_N=BLOCK_N,

431 K_TILES=K_TILES,

424 - STAGES=2,

425 - num_warps=4,

432 + STAGES=STAGES,

433 + num_warps=num_warps,

434 )

435 return logits

436

@@ -521,6 +529,7 @@ def qsa_select_paged_decode(

529 assert token_topk % compress_ratio == 0

530 assert block_indices.shape == (q.shape[0], token_topk // compress_ratio)

531 assert decode_query_len > 0 and q.shape[0] % decode_query_len == 0

532 + assert q.dtype == k_cache.dtype, "Q and the compressed K cache must match"

533 num_requests = q.shape[0] // decode_query_len

534 assert page_table.shape[0] == num_requests

535 assert visible_blocks.shape == (q.shape[0],)

@@ -550,7 +559,8 @@ def qsa_select_paged_decode(

559 BLOCK_N=_DECODE_BLOCK_N,

560 TILES_PER_PROG=tiles_per_program,

561 STAGES=2,

553 - num_warps=2,

562 + # tuned on GB300

563 + num_warps=1 if k_cache.dtype == torch.float8_e4m3fn else 2,

564 )

565 _topk(

566 logits,

@@ -593,6 +603,7 @@ def qsa_select_paged_prefill(

603

604 assert token_topk % compress_ratio == 0

605 assert block_indices.shape == (q.shape[0], token_topk // compress_ratio)

606 + assert q.dtype == k_cache.dtype, "Q and the compressed K cache must match"

607 rows = q.shape[0]

608 # No row scores beyond cdiv(max_seq_len, compress_ratio) compressed

609 # columns. Round up to 64 to keep the logits row stride