[Bugfix][Spec Decode][SM70] Gate DFlash2's BF16 emulation and FlashInfer top-k on the worker's own device

tests/v1/spec_decode/test_dflash2.py

@@ -666,10 +666,15 @@ def test_glm5_dflash_acceptance_policy_preserves_materialize_diagnostic(
 def test_sm70_dflash2_bf16_emulation_has_explicit_ab_switch(monkeypatch):
     config = SimpleNamespace(dtype=torch.bfloat16)
     monkeypatch.setattr(dflash2_model.current_platform, "is_cuda", lambda: True)
+    # The gate asks the worker's own device whether it has SM80; mock a Volta
+    # worker so the switch, not the device query, is under test here.
     monkeypatch.setattr(
         dflash2_model.current_platform,
-        "is_device_capability",
-        lambda capability: capability == 70,
+        "has_device_capability",
+        lambda capability, device_id=0: False,
+    )
+    monkeypatch.setattr(
+        dflash2_model.torch.accelerator, "current_device_index", lambda: 0
     )
     monkeypatch.delenv("VLLM_SM70_DFLASH2_BF16_EMULATION", raising=False)
     assert dflash2_model._use_sm70_bf16_emulation(config)
@@ -1557,7 +1562,10 @@ def test_flashinfer_topk_is_capability_gated_on_sm70(monkeypatch):
     monkeypatch.setattr(
         dflash2_model.current_platform,
         "has_device_capability",
-        lambda capability: capability <= 70,
+        lambda capability, device_id=0: capability <= 70,
+    )
+    monkeypatch.setattr(
+        dflash2_model.torch.accelerator, "current_device_index", lambda: 0
     )
     monkeypatch.setattr(dflash2_model, "has_flashinfer", lambda: True)
     assert dflash2_model._flashinfer_topk() is None

tests/v1/spec_decode/test_dflash2_pre_ampere_gate.py

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DFlash2's BF16 emulation must follow the worker's own device.
+
+``_use_sm70_bf16_emulation`` decides whether a BF16 draft checkpoint runs on
+the range-preserving FP16 path. Every card without native BF16 arithmetic
+needs it -- Volta and Turing alike -- and on a node that mixes generations,
+device 0 of the visibility list belongs to a different worker.
+"""
+
+from dataclasses import dataclass, field
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.model_executor.models import qwen3_dflash2 as dflash2
+
+SM70 = (7, 0)
+SM75 = (7, 5)
+SM80 = (8, 0)
+
+BF16_CONFIG = SimpleNamespace(dtype=torch.bfloat16)
+
+
+@dataclass
+class Node:
+    """Per-index capabilities of a node and the index this worker runs on."""
+
+    capabilities: dict[int, tuple[int, int]] = field(
+        default_factory=lambda: {0: SM75, 1: SM70, 2: SM80}
+    )
+    index: int = 0
+
+
+def _as_tuple(capability: int | tuple[int, int]) -> tuple[int, int]:
+    if isinstance(capability, int):
+        return divmod(capability, 10)
+    return capability
+
+
+@pytest.fixture
+def node(monkeypatch) -> Node:
+    state = Node()
+
+    def fake_has_device_capability(capability, device_id=0):
+        return state.capabilities[device_id] >= _as_tuple(capability)
+
+    def fake_is_device_capability(capability, device_id=0):
+        return state.capabilities[device_id] == _as_tuple(capability)
+
+    monkeypatch.delenv("VLLM_SM70_DFLASH2_BF16_EMULATION", raising=False)
+    monkeypatch.setattr(dflash2.current_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(
+        dflash2.current_platform,
+        "has_device_capability",
+        fake_has_device_capability,
+    )
+    monkeypatch.setattr(
+        dflash2.current_platform,
+        "is_device_capability",
+        fake_is_device_capability,
+    )
+    monkeypatch.setattr(
+        dflash2.torch.accelerator,
+        "current_device_index",
+        lambda: state.index,
+    )
+    return state
+
+
+def test_volta_worker_takes_the_emulation(node: Node) -> None:
+    node.index = 1
+    assert dflash2._use_sm70_bf16_emulation(BF16_CONFIG)
+
+
+def test_turing_worker_takes_the_emulation(node: Node) -> None:
+    """The regression: Turing has no BF16 arithmetic either, but the old
+    exact-SM70 check sent it down the plain FP16 path."""
+    node.index = 0
+    assert dflash2._use_sm70_bf16_emulation(BF16_CONFIG)
+
+
+def test_ampere_worker_keeps_native_bf16(node: Node) -> None:
+    node.index = 2
+    assert not dflash2._use_sm70_bf16_emulation(BF16_CONFIG)
+
+
+def test_gate_does_not_answer_for_device_zero(node: Node) -> None:
+    """Device 0 is Ampere here; the Turing worker must still emulate."""
+    node.capabilities = {0: SM80, 1: SM75}
+    node.index = 1
+    assert dflash2._use_sm70_bf16_emulation(BF16_CONFIG)
+
+
+def test_fp16_checkpoint_needs_no_emulation(node: Node) -> None:
+    node.index = 0
+    fp16_config = SimpleNamespace(dtype=torch.float16)
+    assert not dflash2._use_sm70_bf16_emulation(fp16_config)
+
+
+def test_switch_disables_the_emulation(monkeypatch, node: Node) -> None:
+    monkeypatch.setenv("VLLM_SM70_DFLASH2_BF16_EMULATION", "0")
+    node.index = 0
+    assert not dflash2._use_sm70_bf16_emulation(BF16_CONFIG)
+
+
+def test_non_cuda_platform_never_emulates(monkeypatch) -> None:
+    monkeypatch.setattr(dflash2.current_platform, "is_cuda", lambda: False)
+    assert not dflash2._use_sm70_bf16_emulation(BF16_CONFIG)
+
+
+def test_flashinfer_topk_gate_follows_the_worker_device(
+    monkeypatch, node: Node
+) -> None:
+    """The selector's FlashInfer top-k has no pre-SM80 kernel. Device 0 is
+    Ampere here; the Turing worker must still fall back to torch.topk."""
+    node.capabilities = {0: SM80, 1: SM75}
+    node.index = 1
+    monkeypatch.setattr(dflash2, "has_flashinfer", lambda: True)
+    dflash2._flashinfer_topk.cache_clear()
+    try:
+        assert dflash2._flashinfer_topk() is None
+    finally:
+        dflash2._flashinfer_topk.cache_clear()

vllm/model_executor/models/qwen3_dflash2.py

@@ -57,11 +57,15 @@ def _use_sm70_bf16_emulation(config) -> bool:
     }
     enabled = os.getenv("VLLM_SM70_DFLASH2_BF16_EMULATION", "1").strip().lower()
     enabled = enabled in ("1", "true", "yes", "on")
-    return (
-        enabled
-        and is_bf16
-        and current_platform.is_cuda()
-        and current_platform.is_device_capability(70)
+    if not (enabled and is_bf16 and current_platform.is_cuda()):
+        return False
+    # Native BF16 arithmetic arrives with SM80. Volta and Turing both run the
+    # draft in FP16, so both need the range-preserving path: the criterion is
+    # the missing capability, not one architecture number. Ask the device this
+    # worker builds on; device 0 of the visibility list may be another card on
+    # a node that mixes architectures.
+    return not current_platform.has_device_capability(
+        80, device_id=torch.accelerator.current_device_index()
     )
 
 
@@ -74,7 +78,12 @@ def _flashinfer_topk() -> Callable[..., tuple[torch.Tensor, torch.Tensor]] | Non
     """
     if not current_platform.is_cuda():
         return None
-    if not current_platform.has_device_capability(80):
+    # Same rule as _use_sm70_bf16_emulation: ask the worker's own device, not
+    # device 0 of the visibility list; the cache is per process, i.e. per
+    # worker.
+    if not current_platform.has_device_capability(
+        80, device_id=torch.accelerator.current_device_index()
+    ):
         logger.info_once(
             "DFlash2 disables FlashInfer top-k below SM80; using torch.topk."
         )