From c5a0473f960661b8cb7cc5777bd75a96a5b933cd Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Mon, 20 Jul 2026 15:02:44 -0400 Subject: [PATCH 1/3] glm52: calibrated NVFP4 MLA KV outer scales + serve wiring Per-layer outer-scale calibration for nvfp4_ds_mla KV (VLLM_NVFP4_MLA_SCALES_FILE, format v1) with an explicit default-off knob in serve-glm52.sh. Collapses the NVFP4-vs-FP8 KV KLD gap to ~+0.008-0.009 (0.1345/0.1356 vs 0.1263, 5 fresh boots each, rtx6kpro protocol) while raising max context from 373k to 550k/600k+ on 4x96GB. Co-Authored-By: Claude Fable 5 --- kv-scales/README.md | 49 +++++ ...-nvfp4-nf3-hybrid_mla_outer_scales_v1.json | 178 ++++++++++++++++++ serve-glm52.sh | 2 + 3 files changed, 229 insertions(+) create mode 100644 kv-scales/README.md create mode 100644 kv-scales/glm52-nvfp4-nf3-hybrid_mla_outer_scales_v1.json diff --git a/kv-scales/README.md b/kv-scales/README.md new file mode 100644 index 000000000000..36b37d5aa188 --- /dev/null +++ b/kv-scales/README.md @@ -0,0 +1,49 @@ +# GLM-5.2 NVFP4 MLA KV outer scales + +Per-layer calibration for the `nvfp4_ds_mla` KV cache writer +(`VLLM_NVFP4_MLA_SCALES_FILE`, format `nvfp4_ds_mla_outer_scale_v1`). + +GLM-5.2's post-RMSNorm 512-dim `kv_c` latent spans a ~240x amplitude range +across layers. With the default outer scale of 1.0, shallow layers quantize +with E4M3 block scales at or below the subnormal floor, and shallow-layer KV +error is strongly amplified downstream. `s_l = max_abs(kv_c_normed) / (6*448)` +re-centers every layer so its largest block scale lands at the top of the +E4M3 range. + +## Files + +- `glm52-nvfp4-nf3-hybrid_mla_outer_scales_v1.json` — calibrated on + `madeby561/GLM-5.2-MXFP8-NVFP4-NF3-Hybrid` (K64R16), Salesforce/wikitext + (wikitext-2-raw-v1, test), 2048-token context, TP4; per-layer envelope with + an independent community capture of the same base model for shallow-layer + headroom. `max_abs` per layer is included for auditability. + +## Usage + +```bash +VLLM_NVFP4_MLA_SCALES_FILE=/path/to/glm52-nvfp4-nf3-hybrid_mla_outer_scales_v1.json \ + ./serve-glm52.sh # nvfp4_ds_mla + B12X_MLA_SPARSE only; inert otherwise +``` + +## Results (teacher-forced prefill KLD vs BF16 reference, 5 fresh boots each) + +| KV config | mean +/- sd | max ctx (4x96GB) | +|---|---|---| +| fp8_ds_mla | 0.1263 +/- 0.0030 | 373k | +| nvfp4_ds_mla + scales, bf16 rope | 0.1345 +/- 0.0035 | 550k | +| nvfp4_ds_mla + scales, fp8 rope (`KV_FP8_ROPE=1`) | 0.1356 +/- 0.0054 | 600k+ | +| nvfp4_ds_mla, no scales, bf16 rope | 0.158 | 550k | +| nvfp4_ds_mla, no scales, fp8 rope | 0.168 | 600k+ | + +Protocol: local-inference-lab/rtx6kpro `benchmarks/glm52-kld-evaluation.md` +(festr2 2026-07-08 reference logits, one fixed 2048-token window, 2047 +positions, full 154,880 vocab, `KL(ref || candidate)`). + +## Cache invalidation (required) + +CuTeDSL folds the outer-scale multiply out of the kernel when `latent_scale` +traces at exactly 1.0. b12x builds without the identity/dynamic compile-spec +fact (see lukealonso/b12x PR "mla: split latent_scale identity/dynamic +compile-cache entries") replay stale identity cubins from persistent compile +caches, silently dropping the restore. Clear mounted b12x compile caches once +when enabling scales on such builds. diff --git a/kv-scales/glm52-nvfp4-nf3-hybrid_mla_outer_scales_v1.json b/kv-scales/glm52-nvfp4-nf3-hybrid_mla_outer_scales_v1.json new file mode 100644 index 000000000000..7ae9d9351c4e --- /dev/null +++ b/kv-scales/glm52-nvfp4-nf3-hybrid_mla_outer_scales_v1.json @@ -0,0 +1,178 @@ +{ + "format": "nvfp4_ds_mla_outer_scale_v1", + "num_layers": 78, + "latent_dim": 512, + "denominator": 2688.0, + "formula": "s_l = max_abs(kv_c_normed) / (6 * 448)", + "model": "madeby561/GLM-5.2-MXFP8-NVFP4-NF3-Hybrid (local K64R16 build; own capture)", + "dataset": { + "name": "Salesforce/wikitext", + "config": "wikitext-2-raw-v1", + "split": "test", + "context_length": 2048, + "windows": 1, + "note": "the canonical festr2-0708 census window" + }, + "created_utc": "2026-07-20T16:40:43.312260+00:00", + "hook": "mla.py kv_a_layernorm output (bind-mount capture patch), per-TP-rank agreement checked", + "max_abs": [ + 0.04833984375, + 0.021728515625, + 0.03662109375, + 0.1005859375, + 0.023681640625, + 0.03515625, + 0.046875, + 0.033203125, + 0.259765625, + 0.1474609375, + 0.7734375, + 1.0, + 1.1171875, + 0.150390625, + 0.76171875, + 0.392578125, + 0.25390625, + 0.455078125, + 0.80859375, + 0.80078125, + 0.96875, + 0.78515625, + 0.58203125, + 0.330078125, + 0.89453125, + 0.73046875, + 1.0390625, + 1.578125, + 1.0078125, + 0.92578125, + 1.25, + 1.5859375, + 1.234375, + 1.9140625, + 1.9453125, + 1.7578125, + 1.7890625, + 2.546875, + 2.296875, + 2.4375, + 1.921875, + 2.171875, + 2.875, + 2.546875, + 2.765625, + 3.234375, + 2.5625, + 3.390625, + 2.40625, + 2.546875, + 3.09375, + 2.90625, + 3.671875, + 2.46875, + 2.609375, + 2.359375, + 3.078125, + 3.171875, + 2.578125, + 2.359375, + 3.890625, + 2.953125, + 4.09375, + 4.09375, + 5.0625, + 5.1875, + 2.984375, + 2.875, + 2.984375, + 3.296875, + 3.828125, + 2.859375, + 3.59375, + 3.734375, + 4.25, + 4.1875, + 4.84375, + 3.75 + ], + "scales": [ + 1.7983572823660715e-05, + 8.083525158110119e-06, + 1.3623918805803572e-05, + 3.742036365327381e-05, + 8.810134161086309e-06, + 1.3078962053571428e-05, + 1.743861607142857e-05, + 1.2352353050595238e-05, + 9.663899739583333e-05, + 5.485897972470238e-05, + 0.00028773716517857144, + 0.0003720238095238095, + 0.00041562034970238094, + 5.5948893229166664e-05, + 0.0002833775111607143, + 0.00014604840959821428, + 9.445917038690477e-05, + 0.00016929989769345238, + 0.00030081612723214287, + 0.0002979096912202381, + 0.0003603980654761905, + 0.00029209681919642856, + 0.00021652948288690475, + 0.0001227969215029762, + 0.00033278692336309525, + 0.00027175176711309525, + 0.0003865559895833333, + 0.0005871000744047619, + 0.0003749302455357143, + 0.0003444126674107143, + 0.0004650297619047619, + 0.0005900065104166666, + 0.0004592168898809524, + 0.0007120768229166666, + 0.0007237025669642857, + 0.0006539481026785714, + 0.0006655738467261905, + 0.0009474981398809524, + 0.0008544921875, + 0.0009068080357142857, + 0.0007149832589285714, + 0.0008079892113095238, + 0.0010695684523809525, + 0.0009474981398809524, + 0.0010288783482142857, + 0.0012032645089285715, + 0.0009533110119047619, + 0.0012613932291666667, + 0.0008951822916666666, + 0.0009474981398809524, + 0.0011509486607142857, + 0.0010811941964285715, + 0.001366024925595238, + 0.0009184337797619048, + 0.0009707496279761905, + 0.0008777436755952381, + 0.0011451357886904762, + 0.0011800130208333333, + 0.0009591238839285714, + 0.0008777436755952381, + 0.0014474051339285715, + 0.0010986328125, + 0.0015229724702380952, + 0.0015229724702380952, + 0.0018833705357142857, + 0.001929873511904762, + 0.001110258556547619, + 0.0010695684523809525, + 0.001110258556547619, + 0.0012265159970238095, + 0.0014241536458333333, + 0.0010637555803571428, + 0.0013369605654761905, + 0.0013892764136904762, + 0.0015811011904761905, + 0.0015578497023809525, + 0.0018019903273809525, + 0.0013950892857142857 + ] +} \ No newline at end of file diff --git a/serve-glm52.sh b/serve-glm52.sh index 3af05f4daa9d..e2577952468f 100755 --- a/serve-glm52.sh +++ b/serve-glm52.sh @@ -108,6 +108,8 @@ MOE_BACKEND="${MOE_BACKEND:-b12x}" MOE_SPEC_BACKEND="${MOE_SPEC_BACKEND:-b12x}" ATTENTION_BACKEND="${ATTENTION_BACKEND:-B12X_MLA_SPARSE}" KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-fp8}" +# Calibrated per-layer outer scales for nvfp4_ds_mla KV (empty = disabled). +export VLLM_NVFP4_MLA_SCALES_FILE="${VLLM_NVFP4_MLA_SCALES_FILE:-}" GLM51_PROFILE="${GLM51_PROFILE:-0}" GLM52_CAUSAL_CASCADE="${GLM52_CAUSAL_CASCADE:-0}" GLM52_DSPARK="${GLM52_DSPARK:-0}" From 3e731bc043d23ec21277fb76d3e15fe6da91b23b Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 22 Jul 2026 15:33:48 +0000 Subject: [PATCH 2/3] fix(attention): auto-route B12X MTP verification to decode --- .../v1/attention/test_sparse_mla_backends.py | 220 ++++++++++++++++-- .../attention/backends/mla/b12x_mla_sparse.py | 50 +++- 2 files changed, 248 insertions(+), 22 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 126b5579da70..bc304a4dc03d 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -43,6 +43,7 @@ _global_causal_lens_for_ckv_gather, ) from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( + FlashInferMLASparseSM120Backend, FlashInferMLASparseTRTLLMBackend, ) from vllm.v1.attention.backends.mla.flashmla_sparse import ( @@ -274,6 +275,55 @@ class FakeDCPGroup: ) +@pytest.mark.parametrize( + ("mode", "expected_rows"), + [("0", 2), ("auto", 8), ("1", 16)], +) +def test_b12x_sparse_spec_decode_scratch_capacity( + default_vllm_config, + monkeypatch: pytest.MonkeyPatch, + workspace_init, + mode: str, + expected_rows: int, +) -> None: + if not current_platform.has_device_capability(120): + pytest.skip("B12xMLASparseBackend requires SM 12.0") + if importlib.util.find_spec("sparkinfer") is None: + pytest.skip("sparkinfer package not available") + + default_vllm_config.scheduler_config.max_num_batched_tokens = 64 + default_vllm_config.scheduler_config.max_num_seqs = 2 + default_vllm_config.speculative_config = SimpleNamespace(num_speculative_tokens=3) + monkeypatch.setenv("VLLM_B12X_MLA_SPEC_EXTEND_AS_DECODE", mode) + monkeypatch.setattr( + B12xMLASparseImpl, + "_prewarm_extend_kernels_once", + lambda self, max_batched: None, + ) + + impl = B12xMLASparseImpl( + num_heads=8, + head_size=576, + scale=1.0 / math.sqrt(576), + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="fp8_ds_mla", + logits_soft_cap=None, + attn_type="decoder", + kv_sharing_target_layer_name=None, + topk_indices_buffer=torch.zeros( + (64, 2048), dtype=torch.int32, device=DEVICE_TYPE + ), + kv_lora_rank=512, + qk_nope_head_dim=512, + qk_rope_head_dim=64, + v_head_dim=512, + ) + + assert impl._decode_max_rows == expected_rows + + @pytest.mark.parametrize( ("num_heads", "kernel_num_heads"), [(8, 8), (11, 16), (24, 24)] ) @@ -798,16 +848,16 @@ def test_sparse_backend_decode_correctness( device_capability ): pytest.skip("FlashInferMLASparseTRTLLMBackend requires SM 10.x capability") - elif backend_cls == B12xMLASparseBackend: + elif backend_cls in (B12xMLASparseBackend, FlashInferMLASparseSM120Backend): if not current_platform.has_device_capability(120): - pytest.skip("B12xMLASparseBackend requires SM 12.0 (consumer Blackwell)") - if importlib.util.find_spec("sparkinfer") is None: + pytest.skip(f"{backend_cls.get_name()} requires SM 12.0") + if ( + backend_cls is B12xMLASparseBackend + and importlib.util.find_spec("sparkinfer") is None + ): pytest.skip("sparkinfer package not available") if kv_cache_dtype != "fp8_ds_mla": - # SparkInfer's GLM_NSA kernel consumes the fp8_ds_mla 656 B/token - # record (raw e4m3 + inline FP32 scales). The other cache dtypes are - # advertised for the serving alias path but not exercised here. - pytest.skip("b12x sparse MLA is validated with the fp8_ds_mla cache") + pytest.skip("SM120 sparse MLA is validated with the fp8_ds_mla cache") batch_spec = SPARSE_BACKEND_BATCH_SPECS[batch_name] use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla" @@ -828,7 +878,11 @@ def test_sparse_backend_decode_correctness( # GLM 5.2 selects 2048 rows. Keep the shared backend test small for other # implementations, but validate b12x at the model's real contract instead of # relying on a smaller kernel-supported regime. - topk_tokens = 2048 if backend_cls is B12xMLASparseBackend else 128 + topk_tokens = ( + 2048 + if backend_cls in (B12xMLASparseBackend, FlashInferMLASparseSM120Backend) + else 128 + ) max_seqlen = max(batch_spec.seq_lens) total_cache_tokens = sum(batch_spec.seq_lens) @@ -853,7 +907,11 @@ def test_sparse_backend_decode_correctness( qk_nope_head_dim=qk_nope_head_dim, qk_rope_head_dim=qk_rope_head_dim, v_head_dim=v_head_dim, - model_type="deepseek_v2", + model_type=( + "glm4_moe" + if backend_cls in (B12xMLASparseBackend, FlashInferMLASparseSM120Backend) + else "deepseek_v2" + ), ) model_config.dtype = dtype model_config.get_num_attention_heads = MethodType( @@ -954,16 +1012,18 @@ def test_sparse_backend_decode_correctness( # kernel instead keeps the raw e4m3 K with the inline arbitrary-FP32 # group scale (it is incompatible with ue8m0 block-scaling), so for # B12x the reference must dequantize with the true FP32 scales. - is_sm100 = ( - torch.cuda.get_device_capability()[0] >= 10 - and backend_cls is not B12xMLASparseBackend + uses_pow2_scales = torch.cuda.get_device_capability()[ + 0 + ] >= 10 and backend_cls not in ( + B12xMLASparseBackend, + FlashInferMLASparseSM120Backend, ) kv_c_full, k_pe_squeezed = _quantize_dequantize_fp8_ds_mla( kv_c_full, k_pe_full.squeeze(1), block_size=block_size, scale=kv_cache_scale, - simulate_sm100_e8m0_scales=is_sm100, + simulate_sm100_e8m0_scales=uses_pow2_scales, ) k_pe_full = k_pe_squeezed.unsqueeze(1) @@ -1164,6 +1224,140 @@ def test_sparse_backend_decode_correctness( torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.01, atol=0.01) +@pytest.mark.parametrize( + ("batch_name", "mode", "is_prefilling", "expected_path"), + [ + ("spec_decode_small", "0", False, "extend"), + ("spec_decode_small", "auto", False, "decode"), + ("spec_decode_medium", "auto", False, "decode"), + ("spec_decode_small", "auto", True, "extend"), + ("spec_decode_small", "1", True, "decode"), + ], +) +def test_b12x_sparse_spec_decode_causality( + default_vllm_config, + dist_init, + workspace_init, + monkeypatch: pytest.MonkeyPatch, + batch_name: str, + mode: str, + is_prefilling: bool, + expected_path: str, +) -> None: + """Both verifier paths must match token-wise causal SDPA.""" + if not current_platform.has_device_capability(120): + pytest.skip("B12xMLASparseBackend requires SM 12.0 (consumer Blackwell)") + + sparse_mla = pytest.importorskip("sparkinfer.attention.sparse_mla") + calls = {"decode": 0, "extend": 0} + + def track_path(name, fn): + def wrapped(*args, **kwargs): + calls[name] += 1 + return fn(*args, **kwargs) + + return wrapped + + monkeypatch.setattr( + sparse_mla, + "run_decode", + track_path("decode", sparse_mla.run_decode), + ) + monkeypatch.setattr( + sparse_mla, + "run_extend", + track_path("extend", sparse_mla.run_extend), + ) + monkeypatch.setenv("VLLM_B12X_MLA_SPEC_EXTEND_AS_DECODE", mode) + monkeypatch.setitem(SPARSE_BACKEND_BATCH_SPECS, batch_name, BATCH_SPECS[batch_name]) + + batch_spec = BATCH_SPECS[batch_name] + original_create_vllm_config = create_vllm_config + original_create_common_attn_metadata = create_common_attn_metadata + + def create_spec_vllm_config(*args, **kwargs): + config = original_create_vllm_config(*args, **kwargs) + config.speculative_config = SimpleNamespace( + num_speculative_tokens=max(batch_spec.query_lens) - 1 + ) + return config + + def create_spec_common_attn_metadata(*args, **kwargs): + metadata = original_create_common_attn_metadata(*args, **kwargs) + metadata.is_prefilling = torch.full( + (metadata.num_reqs,), is_prefilling, dtype=torch.bool + ) + return metadata + + monkeypatch.setitem(globals(), "create_vllm_config", create_spec_vllm_config) + monkeypatch.setitem( + globals(), "create_common_attn_metadata", create_spec_common_attn_metadata + ) + + test_sparse_backend_decode_correctness( + default_vllm_config=default_vllm_config, + dist_init=dist_init, + backend_cls=B12xMLASparseBackend, + batch_name=batch_name, + kv_cache_dtype="fp8_ds_mla", + tensor_parallel_size=8, + block_size=64, + workspace_init=workspace_init, + q_scale=1.0, + k_scale=1.0, + ) + + assert calls[expected_path] > 0 + assert calls["decode" if expected_path == "extend" else "extend"] == 0 + + +@pytest.mark.parametrize("batch_name", ["spec_decode_small", "spec_decode_medium"]) +def test_flashinfer_sm120_sparse_spec_decode_causality( + default_vllm_config, + dist_init, + workspace_init, + monkeypatch: pytest.MonkeyPatch, + batch_name: str, +) -> None: + """FlashInfer's flattened verifier rows must remain token-wise causal.""" + if not current_platform.has_device_capability(120): + pytest.skip("FlashInferMLASparseSM120Backend requires SM 12.0") + + from vllm.utils import flashinfer as flashinfer_utils + + calls = [] + original_decode = flashinfer_utils.flashinfer_trtllm_batch_decode_with_kv_cache_mla + + def track_decode(*args, **kwargs): + calls.append(kwargs) + return original_decode(*args, **kwargs) + + monkeypatch.setattr( + flashinfer_utils, + "flashinfer_trtllm_batch_decode_with_kv_cache_mla", + track_decode, + ) + monkeypatch.setitem(SPARSE_BACKEND_BATCH_SPECS, batch_name, BATCH_SPECS[batch_name]) + + test_sparse_backend_decode_correctness( + default_vllm_config=default_vllm_config, + dist_init=dist_init, + backend_cls=FlashInferMLASparseSM120Backend, + batch_name=batch_name, + kv_cache_dtype="fp8_ds_mla", + tensor_parallel_size=8, + block_size=64, + workspace_init=workspace_init, + q_scale=1.0, + k_scale=1.0, + ) + + assert len(calls) == 1 + assert calls[0]["query"].shape[1] == 1 + assert calls[0]["block_tables"].shape[1] == 1 + assert calls[0]["seq_lens"] is None + + def _triton_convert_reference_impl( req_ids: torch.Tensor, block_table: torch.Tensor, diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index f6f8bc0a482c..2e5cea8922cd 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -504,6 +504,9 @@ class B12xMLASparseMetadata(AttentionMetadata): num_decodes: int num_prefills: int prefill_max_seq_len: int + # True only for a multi-token speculative-verification batch. Unlike a + # short chunked prefill, every request has completed its prompt. + is_spec_decode: bool query_start_loc: torch.Tensor slot_mapping: torch.Tensor @@ -565,6 +568,10 @@ def __init__( self.dcp_rank = get_dcp_group().rank_in_group self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size + spec_config = getattr(vllm_config, "speculative_config", None) + self.num_speculative_tokens = int( + getattr(spec_config, "num_speculative_tokens", 0) or 0 + ) max_tokens = vllm_config.scheduler_config.max_num_batched_tokens max_seqs = vllm_config.scheduler_config.max_num_seqs @@ -656,6 +663,14 @@ def build( ) assert num_decode_tokens + num_prefill_tokens == num_tokens + is_spec_decode = False + if ( + self.num_speculative_tokens > 0 + and 1 < cm.max_query_len <= self.num_speculative_tokens + 1 + and cm.is_prefilling is not None + ): + is_spec_decode = not bool(torch.any(cm.is_prefilling[: cm.num_reqs])) + use_dcp = self.dcp_world_size > 1 seq_lens_for_req = ( cm.dcp_local_seq_lens @@ -854,6 +869,7 @@ def build( num_decodes=num_decodes, num_prefills=num_prefills, prefill_max_seq_len=cm.max_seq_len if num_prefills > 0 else 0, + is_spec_decode=is_spec_decode, query_start_loc=cm.query_start_loc, slot_mapping=cm.slot_mapping, block_table=cm.block_table_tensor, @@ -1054,20 +1070,33 @@ def __init__( self._pad_heads = self._kernel_num_heads != self._input_num_heads self.spec_decode_max_q = _env_int("VLLM_B12X_MLA_SPEC_DECODE_MAX_Q", 8) - # The decode kernel handles independent one-token query rows. MTP - # verification has multiple query rows per request, and later rows must - # attend to earlier draft rows in the same verifier batch. Route those - # batches through the extend path unless explicitly overridden. - self.spec_extend_as_decode = ( - os.getenv("VLLM_B12X_MLA_SPEC_EXTEND_AS_DECODE", "0") != "0" + spec_decode_mode = ( + os.getenv("VLLM_B12X_MLA_SPEC_EXTEND_AS_DECODE", "auto").strip().lower() ) + disabled_modes = {"0", "false", "off", "no"} + forced_modes = {"1", "true", "on", "yes"} + if spec_decode_mode not in {"auto", *disabled_modes, *forced_modes}: + raise ValueError( + "VLLM_B12X_MLA_SPEC_EXTEND_AS_DECODE must be auto, 0, or 1 " + f"(got {spec_decode_mode!r})" + ) + self.spec_extend_as_decode = spec_decode_mode not in disabled_modes + self.spec_extend_as_decode_force = spec_decode_mode in forced_modes # Decode query rows per request (1, plus speculative draft tokens). q_per_req = 1 spec = getattr(vllm_config, "speculative_config", None) - if spec is not None and getattr(spec, "num_speculative_tokens", None): + if ( + self.spec_extend_as_decode + and spec is not None + and getattr(spec, "num_speculative_tokens", None) + ): q_per_req = 1 + int(spec.num_speculative_tokens) - if self.spec_extend_as_decode: + # Auto mode only dispatches genuine verifier batches, whose maximum + # row count is fixed by speculative_config. The explicit force mode + # may also route arbitrary short extends and therefore reserves the + # full operator limit. + if self.spec_extend_as_decode_force: q_per_req = max(q_per_req, self.spec_decode_max_q) self._decode_max_rows = min(max_num_seqs * q_per_req, max_batched) if self._decode_max_rows < max_num_seqs: @@ -2189,8 +2218,11 @@ def forward_mqa( B12xMLASparseImpl._shared_gather_event.record(self._ckv_gather_stream) B12xMLASparseImpl._shared_gather_buf_idx = next_buf_idx + use_spec_decode_kernel = self.spec_extend_as_decode and ( + self.spec_extend_as_decode_force or attn_metadata.is_spec_decode + ) use_decode_kernel = attn_metadata.max_query_len <= 1 or ( - self.spec_extend_as_decode + use_spec_decode_kernel and attn_metadata.max_query_len <= self.spec_decode_max_q and num_actual_toks <= attn_metadata.num_reqs * self.spec_decode_max_q and num_actual_toks <= self._decode_max_rows From dc7705900d9420922f5fc150dc63963213014e05 Mon Sep 17 00:00:00 2001 From: derek Date: Thu, 23 Jul 2026 01:52:12 -0400 Subject: [PATCH 3/3] fix(attention): keep compact NVFP4 MTP on qualified path --- .../v1/attention/test_sparse_mla_backends.py | 42 +++++++++++++-- .../attention/backends/mla/b12x_mla_sparse.py | 54 +++++++++++++++---- 2 files changed, 83 insertions(+), 13 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index bc304a4dc03d..594bd3e9efad 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -41,6 +41,7 @@ B12xMLASparseImpl, B12xMLASparseMetadataBuilder, _global_causal_lens_for_ckv_gather, + _resolve_spec_decode_mode, ) from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( FlashInferMLASparseSM120Backend, @@ -81,6 +82,33 @@ DEVICE_TYPE = current_platform.device_type +@pytest.mark.parametrize( + ("mode", "kv_cache_dtype", "expected"), + [ + ("0", "fp8_ds_mla", (False, False)), + ("auto", "fp8_ds_mla", (True, False)), + ("1", "fp8_ds_mla", (True, True)), + ("auto", "nvfp4_ds_mla", (False, False)), + ("1", "nvfp4_ds_mla", (True, True)), + ("auto", "bfloat16", (False, False)), + ], +) +def test_b12x_sparse_spec_decode_auto_is_format_qualified( + mode: str, + kv_cache_dtype: str, + expected: tuple[bool, bool], +) -> None: + assert _resolve_spec_decode_mode( + mode, + kv_cache_dtype=kv_cache_dtype, + ) == expected + + +def test_b12x_sparse_spec_decode_mode_rejects_unknown_value() -> None: + with pytest.raises(ValueError, match="must be auto, 0, or 1"): + _resolve_spec_decode_mode("sometimes", kv_cache_dtype="nvfp4_ds_mla") + + @pytest.mark.parametrize( ("global_seq_lens", "query_start_loc", "req_ids", "expected"), [ @@ -276,14 +304,22 @@ class FakeDCPGroup: @pytest.mark.parametrize( - ("mode", "expected_rows"), - [("0", 2), ("auto", 8), ("1", 16)], + ("mode", "kv_cache_dtype", "expected_rows"), + [ + ("0", "fp8_ds_mla", 2), + ("auto", "fp8_ds_mla", 8), + ("1", "fp8_ds_mla", 16), + ("0", "nvfp4_ds_mla", 2), + ("auto", "nvfp4_ds_mla", 2), + ("1", "nvfp4_ds_mla", 16), + ], ) def test_b12x_sparse_spec_decode_scratch_capacity( default_vllm_config, monkeypatch: pytest.MonkeyPatch, workspace_init, mode: str, + kv_cache_dtype: str, expected_rows: int, ) -> None: if not current_platform.has_device_capability(120): @@ -308,7 +344,7 @@ def test_b12x_sparse_spec_decode_scratch_capacity( num_kv_heads=1, alibi_slopes=None, sliding_window=None, - kv_cache_dtype="fp8_ds_mla", + kv_cache_dtype=kv_cache_dtype, logits_soft_cap=None, attn_type="decoder", kv_sharing_target_layer_name=None, diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index 2e5cea8922cd..8c20f667a3a4 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -146,6 +146,32 @@ def _env_int(name: str, default: int) -> int: return parsed +def _resolve_spec_decode_mode( + raw_mode: str, *, kv_cache_dtype: str +) -> tuple[bool, bool]: + """Resolve the MTP verifier route without overextending ``auto``. + + The flattened decode-verifier path is numerically qualified by + ``test_b12x_sparse_spec_decode_causality`` only for ``fp8_ds_mla``. The + compact ``nvfp4_ds_mla`` format uses a different BF16-QK kernel arm and was + never covered by that qualification, so auto mode keeps it on the established + extend path. Explicit force mode remains available for focused bring-up. + + Returns ``(use_decode_for_verifier, explicitly_forced)``. + """ + mode = raw_mode.strip().lower() + disabled_modes = {"0", "false", "off", "no"} + forced_modes = {"1", "true", "on", "yes"} + if mode not in {"auto", *disabled_modes, *forced_modes}: + raise ValueError( + "VLLM_B12X_MLA_SPEC_EXTEND_AS_DECODE must be auto, 0, or 1 " + f"(got {mode!r})" + ) + if mode == "auto": + return kv_cache_dtype == "fp8_ds_mla", False + return mode in forced_modes, mode in forced_modes + + def _get_ckv_gather_workspace(device: torch.device, nbytes: int) -> torch.Tensor: key = (device.type, device.index) workspace = _CKV_GATHER_WORKSPACES.get(key) @@ -1070,18 +1096,26 @@ def __init__( self._pad_heads = self._kernel_num_heads != self._input_num_heads self.spec_decode_max_q = _env_int("VLLM_B12X_MLA_SPEC_DECODE_MAX_Q", 8) - spec_decode_mode = ( - os.getenv("VLLM_B12X_MLA_SPEC_EXTEND_AS_DECODE", "auto").strip().lower() + spec_decode_mode = os.getenv( + "VLLM_B12X_MLA_SPEC_EXTEND_AS_DECODE", "auto" ) - disabled_modes = {"0", "false", "off", "no"} - forced_modes = {"1", "true", "on", "yes"} - if spec_decode_mode not in {"auto", *disabled_modes, *forced_modes}: - raise ValueError( - "VLLM_B12X_MLA_SPEC_EXTEND_AS_DECODE must be auto, 0, or 1 " - f"(got {spec_decode_mode!r})" + ( + self.spec_extend_as_decode, + self.spec_extend_as_decode_force, + ) = _resolve_spec_decode_mode( + spec_decode_mode, + kv_cache_dtype=self.kv_cache_dtype, + ) + if ( + spec_decode_mode.strip().lower() == "auto" + and not self.spec_extend_as_decode + ): + logger.info_once( + "B12X MTP verifier auto-route keeps kv_cache_dtype=%s on the " + "numerically qualified extend path; set " + "VLLM_B12X_MLA_SPEC_EXTEND_AS_DECODE=1 to force decode", + self.kv_cache_dtype, ) - self.spec_extend_as_decode = spec_decode_mode not in disabled_modes - self.spec_extend_as_decode_force = spec_decode_mode in forced_modes # Decode query rows per request (1, plus speculative draft tokens). q_per_req = 1