diff --git a/docs/design/flashnext_16gb_host_kv.md b/docs/design/flashnext_16gb_host_kv.md index 0f65481f5..4387786ca 100644 --- a/docs/design/flashnext_16gb_host_kv.md +++ b/docs/design/flashnext_16gb_host_kv.md @@ -321,3 +321,10 @@ Ordinary dense segment scratch is bounded by the admitted K2560/N4096 shapes, and shared scratch by K2560/N320 or K160/N2560. Code, coefficient and partial storage decreases by 34.5 MiB per device without changing their layouts, arithmetic or supported shapes. + +The subsequent M5 graph capture reached the fused side projection and failed +because its FP32 split-K workspace inherited the model loader's FP16 default +dtype. Allocate that workspace explicitly as FP32. Six formats passed both +FP16 and FP32 default-dtype tests, shared-bank checks and changed-input M5/M20 +graphs (12 tests). This repairs workspace typing; it does not change weight +or activation precision. Full-model serving remains pending. diff --git a/tests/kernels/quantization/test_sm70_dmv13_side.py b/tests/kernels/quantization/test_sm70_dmv13_side.py index 653d2392c..3b3d5e083 100644 --- a/tests/kernels/quantization/test_sm70_dmv13_side.py +++ b/tests/kernels/quantization/test_sm70_dmv13_side.py @@ -16,6 +16,16 @@ pytestmark = pytest.mark.skipif( BLOCK = {8: 34, 12: 144, 13: 176, 14: 210, 20: 18, 23: 136} +@pytest.fixture(params=[torch.float32, torch.float16]) +def default_dtype(request): + previous = torch.get_default_dtype() + torch.set_default_dtype(request.param) + try: + yield + finally: + torch.set_default_dtype(previous) + + def _raw(rng, rows, k, qtype): size = BLOCK[qtype] block = 32 if qtype in (8, 20) else 256 @@ -84,7 +94,7 @@ def test_side_projection_matches_dequant(qtype, tokens, merged_bf16): @pytest.mark.parametrize("qtype", [8, 12, 13, 14, 20, 23]) -def test_shared_segment_planes_preserve_side_and_m20_graph(qtype): +def test_shared_segment_planes_preserve_side_and_m20_graph(qtype, default_dtype): from types import SimpleNamespace from vllm.model_executor.layers.quantization.gguf_dense_hmma_formats import ( @@ -113,6 +123,7 @@ def test_shared_segment_planes_preserve_side_and_m20_graph(qtype): extra = (torch.randn(24, k, device="cuda") * 0.02).half() proj = Dmv13Projection(_Shard(raw, qtype), _Extra(extra)) assert proj.ready + assert proj.workspace.dtype == torch.float32 layer.sm70_side_projection = proj _register_banks(layer, "sm70_side_projection", proj) x = torch.randn(5, k, device="cuda", dtype=torch.float16) diff --git a/vllm/model_executor/layers/quantization/sm70_dmv13_projection.py b/vllm/model_executor/layers/quantization/sm70_dmv13_projection.py index 40cfcdc7e..cb6037f48 100644 --- a/vllm/model_executor/layers/quantization/sm70_dmv13_projection.py +++ b/vllm/model_executor/layers/quantization/sm70_dmv13_projection.py @@ -121,7 +121,7 @@ class Dmv13Projection: return self.extra = weight.contiguous() tiles = sum((w + 31) // 32 for w in self.widths) - self.workspace = torch.zeros(tiles * 256, device=device) + self.workspace = torch.zeros(tiles * 256, device=device, dtype=torch.float32) self.counters = torch.zeros(tiles, device=device, dtype=torch.int32) self.ready = hasattr(torch.ops._C, "sm70_dmv13_out")