[Bugfix] Generalize PLE graph and runtime compatibility

flash-attention-v100/include/fused_mha.h

@@ -59,6 +59,8 @@ at::Tensor flash_attention_grouped_verify_paged(
 
 int64_t flash_attention_grouped_verify_max_query_tokens();
 
+int64_t flash_attention_grouped_sparse_page4_abi_version();
+
 at::Tensor flash_attention_grouped_sparse_page4(
     const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache,
     std::optional<at::Tensor>& out_, const at::Tensor& block_table,

flash-attention-v100/kernel/flash_decode_paged.cu

@@ -4034,6 +4034,12 @@ int64_t flash_attention_grouped_verify_max_query_tokens() {
   return kGroupedVerifyMaxSupportedQ;
 }
 
+int64_t flash_attention_grouped_sparse_page4_abi_version() {
+  // Version 1 accepted FP16 K/V through the nine-argument forward binding.
+  // Version 2 adds kv_cache_dtype and calibrated K/V scales.
+  return 2;
+}
+
 at::Tensor flash_attention_grouped_verify_paged(
     const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache,
     std::optional<at::Tensor>& out_, const at::Tensor& block_table,

flash-attention-v100/kernel/fused_mha_api.cpp

@@ -29,6 +29,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         "Maximum query length supported by grouped DFlash2 verification");
   m.def("grouped_sparse_page4_fwd", &flash_attention_grouped_sparse_page4,
         "Grouped exact QSA page4 attention over paged KV cache (Volta)");
+  m.def("grouped_sparse_page4_abi_version",
+        &flash_attention_grouped_sparse_page4_abi_version,
+        "Grouped sparse page4 forward ABI version");
   m.def("grouped_sparse_page4_plan_fwd",
         &flash_attention_grouped_sparse_page4_plan,
         "Build grouped exact QSA page4 tables over paged KV cache (Volta)");

tests/engine/test_arg_utils.py

@@ -308,19 +308,8 @@ def test_sm70_mtp_defaults_require_env_opt_in(monkeypatch):
     }
     assert args.enable_prefix_caching is True
     assert args.mamba_cache_mode == "align"
-    assert args.max_num_seqs == 4
-    assert args.compilation_config.cudagraph_capture_sizes == [
-        1,
-        2,
-        4,
-        5,
-        8,
-        9,
-        10,
-        15,
-        18,
-        20,
-    ]
+    assert args.max_num_seqs is None
+    assert args.compilation_config.cudagraph_capture_sizes is None
 
 
 def test_sm70_mtp_split_cudagraphs_are_opt_in(monkeypatch):
@@ -332,7 +321,7 @@ def test_sm70_mtp_split_cudagraphs_are_opt_in(monkeypatch):
         },
     )
 
-    assert args.compilation_config.cudagraph_capture_sizes == [5, 10, 20]
+    assert args.compilation_config.cudagraph_capture_sizes is None
 
 
 def test_sm70_mtp_split_cudagraphs_cover_production_batches(monkeypatch):
@@ -348,6 +337,7 @@ def test_sm70_mtp_split_cudagraphs_cover_production_batches(monkeypatch):
     assert args.compilation_config.cudagraph_capture_sizes == [
         5,
         10,
+        15,
         20,
         30,
         40,
@@ -395,7 +385,7 @@ def test_sm70_explicit_mtp_still_gets_safe_defaults(monkeypatch):
     }
     assert args.enable_prefix_caching is True
     assert args.mamba_cache_mode == "align"
-    assert args.max_num_seqs == 4
+    assert args.max_num_seqs is None
 
 
 def test_sm70_explicit_dflash_preserves_probabilistic_default(monkeypatch):

tests/kernels/attention/test_sm70_qsa_grouped_page4.py

@@ -21,10 +21,21 @@ def _require_grouped_page4():
     return extension
 
 
+def _grouped_page4_abi_version(extension) -> int:
+    capability = getattr(extension, "grouped_sparse_page4_abi_version", None)
+    if callable(capability):
+        return int(capability())
+    doc = getattr(extension.grouped_sparse_page4_fwd, "__doc__", "") or ""
+    return 2 if "arg11:" in doc else 1 if "arg8:" in doc else 0
+
+
 @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3"])
 @torch.inference_mode()
 def test_sm70_qsa_grouped_page4_calibrated_kv(kv_cache_dtype: str) -> None:
     extension = _require_grouped_page4()
+    abi_version = _grouped_page4_abi_version(extension)
+    if kv_cache_dtype == "fp8_e4m3" and abi_version < 2:
+        pytest.skip("installed grouped page4 ABI does not support quantized K/V")
     torch.manual_seed(7)
     query = torch.randn((8, 6, 256), dtype=torch.float16, device="cuda") * 0.2
     key = torch.randn((1, 4, 1, 256), dtype=torch.float16, device="cuda") * 0.35
@@ -48,7 +59,7 @@ def test_sm70_qsa_grouped_page4_calibrated_kv(kv_cache_dtype: str) -> None:
     seq_lens = torch.tensor([4], dtype=torch.int32, device="cuda")
     output = torch.empty_like(query)
     lse = torch.empty((8, 6), dtype=torch.float32, device="cuda")
-    extension.grouped_sparse_page4_fwd(
+    args = (
         query,
         key_cache,
         value_cache,
@@ -58,10 +69,16 @@ def test_sm70_qsa_grouped_page4_calibrated_kv(kv_cache_dtype: str) -> None:
         seq_lens,
         lse,
         256**-0.5,
-        kv_cache_dtype,
-        k_scale,
-        v_scale,
     )
+    if abi_version >= 2:
+        extension.grouped_sparse_page4_fwd(
+            *args,
+            kv_cache_dtype,
+            k_scale,
+            v_scale,
+        )
+    else:
+        extension.grouped_sparse_page4_fwd(*args)
 
     reference_key = reference_key.view(4, 256)
     reference_value = reference_value.view(4, 256)

tests/kernels/moe/test_sm70_unquantized_moe_config.py

@@ -8,6 +8,7 @@
 from vllm.model_executor.layers.fused_moe.fused_moe import (
     _get_sm70_mtp_moe_decode_config,
     force_sm70_mtp_moe_legacy_config,
+    fused_moe_kernel,
 )
 
 
@@ -20,6 +21,10 @@ def test_mtp_sm70_decode_config_keeps_legacy_tile_at_m1():
     assert _get_sm70_mtp_moe_decode_config(1, 256, 128, 2048, 8) is None
 
 
+def test_fused_moe_does_not_specialize_on_routing_dependent_em_alignment():
+    assert "EM" in fused_moe_kernel.do_not_specialize_on_alignment
+
+
 @pytest.mark.parametrize("m", range(2, 17))
 def test_mtp_sm70_decode_config_uses_exact_local_tile(m):
     config = _get_sm70_mtp_moe_decode_config(m, 256, 128, 2048, 8)

tests/models/qwen4_exp/test_qsa_ops.py

@@ -162,6 +162,66 @@ def test_qsa_xqa_page4_route_uses_configured_boundary(monkeypatch):
     assert not qsa_ops._use_sm70_qsa_xqa_page4(query, *args)
 
 
+def test_qsa_grouped_page4_modern_abi_forwards_quantized_kv_metadata():
+    calls = []
+
+    def forward(*args):
+        calls.append(args)
+
+    extension = SimpleNamespace(
+        grouped_sparse_page4_abi_version=lambda: 2,
+        grouped_sparse_page4_plan_fwd=lambda *args: None,
+        grouped_sparse_page4_fwd=forward,
+    )
+    tensors = [torch.empty(0) for _ in range(8)]
+
+    assert qsa_ops._qsa_grouped_page4_supported(extension, "auto")
+    assert qsa_ops._qsa_grouped_page4_supported(extension, "fp8_e4m3")
+    qsa_ops._qsa_grouped_page4_forward(
+        extension,
+        *tensors,
+        0.0625,
+        "fp8_e4m3",
+        0.125,
+        0.25,
+    )
+
+    assert len(calls) == 1
+    assert len(calls[0]) == 12
+    assert calls[0][-3:] == ("fp8_e4m3", 0.125, 0.25)
+
+
+def test_qsa_grouped_page4_legacy_abi_is_fp16_only():
+    calls = []
+
+    def forward(*args):
+        calls.append(args)
+
+    forward.__doc__ = "grouped_sparse_page4_fwd(" + ", ".join(
+        f"arg{index}: object" for index in range(9)
+    )
+    extension = SimpleNamespace(
+        grouped_sparse_page4_plan_fwd=lambda *args: None,
+        grouped_sparse_page4_fwd=forward,
+    )
+    tensors = [torch.empty(0) for _ in range(8)]
+
+    assert qsa_ops._qsa_grouped_page4_abi_version(extension) == 1
+    assert qsa_ops._qsa_grouped_page4_supported(extension, "auto")
+    assert not qsa_ops._qsa_grouped_page4_supported(extension, "fp8_e4m3")
+    qsa_ops._qsa_grouped_page4_forward(
+        extension,
+        *tensors,
+        0.0625,
+        "auto",
+        1.0,
+        1.0,
+    )
+
+    assert len(calls) == 1
+    assert len(calls[0]) == 9
+
+
 def test_qsa_e4m3_page4_routes_large_mixed_batch_below_prefill_boundary(
     monkeypatch,
 ):
@@ -194,13 +254,21 @@ def test_qsa_e4m3_page4_routes_large_mixed_batch_below_prefill_boundary(
     )
 
 
-def test_qsa_e4m3_xqa_page4_splits_non_grouped_large_batch(monkeypatch):
-    rows = 49
+@pytest.mark.parametrize(
+    ("rows", "kv_cache_dtype"),
+    [(49, "fp8_e4m3"), (65, "auto")],
+)
+def test_qsa_xqa_page4_splits_non_grouped_large_batch(
+    monkeypatch,
+    rows,
+    kv_cache_dtype,
+):
     query = torch.empty(rows, 6, 256, dtype=torch.float16)
     flash_cuda = SimpleNamespace(
         decode_paged_xqa_fwd=object(),
-        grouped_sparse_page4_plan_fwd=object(),
-        grouped_sparse_page4_fwd=object(),
+        grouped_sparse_page4_abi_version=lambda: 2,
+        grouped_sparse_page4_plan_fwd=lambda *args: None,
+        grouped_sparse_page4_fwd=lambda *args: None,
     )
     flash_interface = ModuleType("flash_attn_v100.flash_attn_interface")
     cast(Any, flash_interface).flash_attn_v100_cuda = flash_cuda
@@ -292,15 +360,23 @@ def fake_xqa_batch(
         query_positions,
         sequence_lengths,
         out,
-        "fp8_e4m3",
+        kv_cache_dtype,
         0.05,
         0.05,
     )
 
     assert result is out
+    grouped_rows = rows // 8 * 8
     assert calls == [
-        ("grouped", 48, 48, 48, 48, 48),
-        ("xqa", 1, 1, 1, 1, 1),
+        (
+            "grouped",
+            grouped_rows,
+            grouped_rows,
+            grouped_rows,
+            grouped_rows,
+            grouped_rows,
+        ),
+        ("xqa", *(rows - grouped_rows,) * 5),
     ]
 
 

tests/quantization/test_sm70_modelopt_mixed_nvfp4.py

@@ -21,13 +21,23 @@
     _prepare_compact_slot_groups,
     _prepare_single_token_slots,
     _single_token_weighted_reduce,
+    _use_compact_grouped,
     _use_qwen38_indexed_prefill,
     _use_qwen38_qpn_m1_decode,
     _validate_weight_layout,
     validate_nvfp4_sm70_moe_contract,
 )
 
 
+@pytest.mark.parametrize(
+    ("top_k", "compact_tokens", "dense_tokens"),
+    [(8, 10, 11), (10, 8, 9)],
+)
+def test_nvfp4_compact_work_limit_is_routed_rows(top_k, compact_tokens, dense_tokens):
+    assert _use_compact_grouped(compact_tokens, top_k)
+    assert not _use_compact_grouped(dense_tokens, top_k)
+
+
 def _mixed_config() -> ModelOptMixedPrecisionConfig:
     fp8 = ModelOptFp8Config("FP8", True, None, [])
     nvfp4 = ModelOptNvFp4Config(
@@ -249,7 +259,7 @@ def test_nvfp4_sm70_moe_owns_routing_without_generic_modular_wrapper():
     not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0),
     reason="requires an exact SM70 CUDA device",
 )
-@pytest.mark.parametrize("total_slots", (8, 72, 80, 100))
+@pytest.mark.parametrize("total_slots", (8, 72, 80))
 def test_nvfp4_compact_groups_keep_duplicate_expert_slots_independent(total_slots):
     sorted_expert_ids = (
         torch.arange(total_slots, dtype=torch.int32, device="cuda") // 3
@@ -266,6 +276,18 @@ def test_nvfp4_compact_groups_keep_duplicate_expert_slots_independent(total_slot
     assert torch.equal(active_expert_ids.cpu(), sorted_expert_ids.cpu())
 
 
+@pytest.mark.parametrize("total_slots", (81, 100))
+def test_nvfp4_compact_groups_reject_work_above_80_rows(total_slots):
+    sorted_expert_ids = torch.empty(total_slots, dtype=torch.int32)
+    compact_offsets = torch.empty(total_slots + 1, dtype=torch.int32)
+    active_expert_ids = torch.empty(total_slots, dtype=torch.int32)
+
+    with pytest.raises(ValueError, match="active-expert slots"):
+        _prepare_compact_slot_groups(
+            sorted_expert_ids, compact_offsets, active_expert_ids
+        )
+
+
 @pytest.mark.skipif(
     not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0),
     reason="requires an exact SM70 CUDA device",

tests/quantization/test_sm70_warmup.py

@@ -278,7 +278,7 @@ def _nvfp4_moe_layer() -> nn.Module:
     layer.sm70_nvfp4_w2_n_dim = 2048
     layer.sm70_nvfp4_group_size = 16
     layer.sm70_nvfp4_graph_safe_max_tokens = 18
-    layer.sm70_nvfp4_compact_grouped_max_tokens = 10
+    layer.sm70_nvfp4_compact_grouped_max_slots = 80
     layer.w13_tm_weight = nn.Parameter(
         torch.empty((1, 1), dtype=torch.uint8), requires_grad=False
     )
@@ -366,7 +366,7 @@ def test_nvfp4_moe_warmup_includes_opted_in_cuda_graph_shapes(monkeypatch):
     ) == [*range(1, 11), 18, 20, 40, 60, 80]
 
 
-def test_nvfp4_moe_warmup_uses_slot_compact_through_b10(monkeypatch):
+def test_nvfp4_moe_warmup_uses_slot_compact_through_80_rows(monkeypatch):
     layer = _nvfp4_moe_layer()
     calls = []
     monkeypatch.setattr(
@@ -400,7 +400,7 @@ def test_nvfp4_moe_warmup_uses_slot_compact_through_b10(monkeypatch):
     assert calls[2][2].tolist() == list(range(81))
 
 
-def test_nvfp4_moe_warmup_uses_full_expert_groups_above_compact_b10(monkeypatch):
+def test_nvfp4_moe_warmup_uses_full_expert_groups_above_80_rows(monkeypatch):
     layer = _nvfp4_moe_layer()
     calls = []
     monkeypatch.setattr(

tests/v1/attention/test_sm70_flash_v100_policy.py

@@ -1581,11 +1581,11 @@ def test_sm70_nomtp_cudagraph_capture_sizes_cover_concurrency(
     [
         (1, [5]),
         (2, [5, 10]),
-        (4, [5, 10, 20]),
-        (6, [5, 10, 20, 30]),
-        (12, [5, 10, 20, 30, 40, 60]),
-        (16, [5, 10, 20, 30, 40, 60, 80]),
-        (32, [5, 10, 20, 30, 40, 60, 80]),
+        (4, [5, 10, 15, 20]),
+        (6, [5, 10, 15, 20, 30]),
+        (12, [5, 10, 15, 20, 30, 40, 60]),
+        (16, [5, 10, 15, 20, 30, 40, 60, 80]),
+        (32, [5, 10, 15, 20, 30, 40, 60, 80]),
     ],
 )
 def test_sm70_mtp_cudagraph_capture_sizes_cover_production_concurrency(
@@ -1597,6 +1597,24 @@ def test_sm70_mtp_cudagraph_capture_sizes_cover_production_concurrency(
     assert _sm70_mtp_cudagraph_capture_sizes(max_num_seqs, 5) == expected
 
 
+def test_sm70_speculative_cudagraph_shapes_are_tp_independent_and_bounded():
+    from vllm.config.vllm import _sm70_speculative_cudagraph_capture_sizes
+
+    assert _sm70_speculative_cudagraph_capture_sizes(4, 5) == [
+        1,
+        2,
+        4,
+        5,
+        8,
+        9,
+        10,
+        15,
+        18,
+        20,
+    ]
+    assert _sm70_speculative_cudagraph_capture_sizes(256, 5)[-1] == 80
+
+
 def test_flash_v100_decode_query_does_not_attach_smallq_metadata(
     monkeypatch,
     local_flash_v100_model,