diff --git a/driver/cuda/src/model/qwen3_5.cpp b/driver/cuda/src/model/qwen3_5.cpp index 807aac2b9..af6d5ce19 100644 --- a/driver/cuda/src/model/qwen3_5.cpp +++ b/driver/cuda/src/model/qwen3_5.cpp @@ -320,6 +320,12 @@ Qwen3_5Weights bind_qwen3_5(LoadedModel& engine) { const int rank = engine.distributed().tp_rank; const int K_dim = cfg.linear_num_key_heads * cfg.linear_key_head_dim; const int V_dim = cfg.linear_num_value_heads * cfg.linear_value_head_dim; + const auto qkv_quant = engine.quant_meta(la + "in_proj_qkv.weight"); + if (T > 1 && qkv_quant.has_value()) { + throw std::runtime_error( + "qwen3_5: TP-sliced FP8 linear_attn.in_proj_qkv " + "requires matching scale slicing"); + } if (T > 1) { w.owned_bf16_buffers.push_back( slice_la_kkv_blocked(*full_qkv, K_dim, V_dim, rank, T)); @@ -342,7 +348,17 @@ Qwen3_5Weights bind_qwen3_5(LoadedModel& engine) { Lw.la_in_proj_z = &must(engine, la + "in_proj_z.weight"); Lw.la_in_proj_b = &must(engine, la + "in_proj_b.weight"); Lw.la_in_proj_a = &must(engine, la + "in_proj_a.weight"); - if (fused_gdn_projection_weights_enabled()) { + Lw.la_in_proj_qkv_quant = qkv_quant; + Lw.la_in_proj_z_quant = engine.quant_meta(la + "in_proj_z.weight"); + Lw.la_in_proj_b_quant = engine.quant_meta(la + "in_proj_b.weight"); + Lw.la_in_proj_a_quant = engine.quant_meta(la + "in_proj_a.weight"); + const bool can_fuse_gdn_projection_weights = + fused_gdn_projection_weights_enabled() && + !Lw.la_in_proj_qkv_quant.has_value() && + !Lw.la_in_proj_z_quant.has_value() && + !Lw.la_in_proj_b_quant.has_value() && + !Lw.la_in_proj_a_quant.has_value(); + if (can_fuse_gdn_projection_weights) { w.owned_bf16_buffers.push_back(concat_axis0_bf16( *Lw.la_in_proj_qkv, *Lw.la_in_proj_z, "qwen3_5: fuse linear_attn.in_proj_qkvz")); @@ -365,6 +381,7 @@ Qwen3_5Weights bind_qwen3_5(LoadedModel& engine) { w.owned_fp32_buffers.push_back(to_fp32(must(engine, la + "norm.weight"))); Lw.la_norm_w_fp32 = w.owned_fp32_buffers.back().data(); Lw.la_out_proj = &must(engine, la + "out_proj.weight"); + Lw.la_out_proj_quant = engine.quant_meta(la + "out_proj.weight"); Lw.kv_layer = -1; } else if (kind == "full_attention") { Lw.kind = Qwen3_5LayerWeights::Kind::FullAttn; diff --git a/driver/cuda/src/model/qwen3_5.hpp b/driver/cuda/src/model/qwen3_5.hpp index 2c2a3a031..7eb4ec253 100644 --- a/driver/cuda/src/model/qwen3_5.hpp +++ b/driver/cuda/src/model/qwen3_5.hpp @@ -77,8 +77,11 @@ struct Qwen3_5LayerWeights { const DeviceTensor* gate_up_proj_fused = nullptr; // [2*I, H] bf16 // Optional QuantMeta companions for the GEMM-fed projections. The // materialized WeightStore owns this metadata after storage-program execution. - // Linear-attn weights stay bf16 for now (their fused [K1|K2|V] block - // layout needs per-block scale handling that isn't wired yet). + std::optional la_in_proj_qkv_quant; + std::optional la_in_proj_z_quant; + std::optional la_in_proj_b_quant; + std::optional la_in_proj_a_quant; + std::optional la_out_proj_quant; std::optional fa_q_proj_quant; std::optional fa_k_proj_quant; std::optional fa_v_proj_quant; diff --git a/driver/cuda/src/model/qwen3_5_forward.cpp b/driver/cuda/src/model/qwen3_5_forward.cpp index 40f775ad9..c4fc881f9 100644 --- a/driver/cuda/src/model/qwen3_5_forward.cpp +++ b/driver/cuda/src/model/qwen3_5_forward.cpp @@ -538,18 +538,22 @@ void linear_attn_layer_body( } else { // mixed_qkv [N, conv_dim] = norm_x @ in_proj_qkv.T ops::gemm_act_x_w(cublas.handle(), - ws.norm_x.data(), *Lw.la_in_proj_qkv, + ws.norm_x.data(), + make_weight_view(Lw.la_in_proj_qkv, Lw.la_in_proj_qkv_quant), la.mixed_qkv.data(), N, conv_dim, H); // z [N, V_dim] = norm_x @ in_proj_z.T ops::gemm_act_x_w(cublas.handle(), - ws.norm_x.data(), *Lw.la_in_proj_z, + ws.norm_x.data(), + make_weight_view(Lw.la_in_proj_z, Lw.la_in_proj_z_quant), la.z.data(), N, V_dim, H); // a [N, V_h] = norm_x @ in_proj_a.T (b symmetric) ops::gemm_act_x_w(cublas.handle(), - ws.norm_x.data(), *Lw.la_in_proj_a, + ws.norm_x.data(), + make_weight_view(Lw.la_in_proj_a, Lw.la_in_proj_a_quant), la.a.data(), N, V_h, H); ops::gemm_act_x_w(cublas.handle(), - ws.norm_x.data(), *Lw.la_in_proj_b, + ws.norm_x.data(), + make_weight_view(Lw.la_in_proj_b, Lw.la_in_proj_b_quant), la.b.data(), N, V_h, H); } if (stash_write) { @@ -954,11 +958,13 @@ void linear_attn_layer_body( profile_forward_stage_ptr(profile, &ForwardProfile::linear_out_ms, stream, [&] { if (T == 1) { ops::gemm_act_x_w(cublas.handle(), - la.core_out_bf16.data(), *Lw.la_out_proj, + la.core_out_bf16.data(), + make_weight_view(Lw.la_out_proj, Lw.la_out_proj_quant), ws.y.data(), N, H, V_dim, /*beta=*/1.f); } else { ops::gemm_act_x_w(cublas.handle(), - la.core_out_bf16.data(), *Lw.la_out_proj, + la.core_out_bf16.data(), + make_weight_view(Lw.la_out_proj, Lw.la_out_proj_quant), ws.norm_y.data(), N, H, V_dim, /*beta=*/0.f); tp->all_reduce_bf16(ws.norm_y.data(), static_cast(N) * H, ncclSum, stream); diff --git a/driver/cuda/tests/test_qwen35_linear_attn_quantmeta.py b/driver/cuda/tests/test_qwen35_linear_attn_quantmeta.py new file mode 100644 index 000000000..b2bbcc7b0 --- /dev/null +++ b/driver/cuda/tests/test_qwen35_linear_attn_quantmeta.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +"""Regression checks for dense Qwen3.5 linear-attention FP8 metadata wiring.""" + +from pathlib import Path +import unittest + + +ROOT = Path(__file__).resolve().parents[1] +MODEL_DIR = ROOT / "src" / "model" + + +def read_model_file(name: str) -> str: + return (MODEL_DIR / name).read_text() + + +class Qwen35LinearAttnQuantMetaTests(unittest.TestCase): + def test_qwen35_linear_attention_weights_keep_quantmeta_companions(self) -> None: + header = read_model_file("qwen3_5.hpp") + + for field in ( + "la_in_proj_qkv_quant", + "la_in_proj_z_quant", + "la_in_proj_a_quant", + "la_in_proj_b_quant", + "la_out_proj_quant", + ): + self.assertIn(f"std::optional {field};", header) + + def test_qwen35_bind_populates_linear_attention_quantmeta_from_weight_store( + self, + ) -> None: + binder = read_model_file("qwen3_5.cpp") + + self.assertRegex( + binder, + r'const auto qkv_quant\s*=\s*engine\.quant_meta\(la \+ "in_proj_qkv.weight"\);', + ) + self.assertIn("Lw.la_in_proj_qkv_quant = qkv_quant;", binder) + + for field, weight_name in ( + ("la_in_proj_z_quant", "in_proj_z.weight"), + ("la_in_proj_a_quant", "in_proj_a.weight"), + ("la_in_proj_b_quant", "in_proj_b.weight"), + ("la_out_proj_quant", "out_proj.weight"), + ): + self.assertRegex( + binder, + rf'Lw\.{field}\s*=\s*engine\.quant_meta\(la \+ "{weight_name}"\);', + ) + + self.assertIn("const bool can_fuse_gdn_projection_weights =", binder) + for field in ( + "la_in_proj_qkv_quant", + "la_in_proj_z_quant", + "la_in_proj_a_quant", + "la_in_proj_b_quant", + ): + self.assertIn(f"!Lw.{field}.has_value()", binder) + + def test_qwen35_bind_rejects_fp8_qkv_when_tp_slice_would_drop_scale_slice( + self, + ) -> None: + binder = read_model_file("qwen3_5.cpp") + + self.assertIn("T > 1 && qkv_quant.has_value()", binder) + self.assertIn("TP-sliced FP8 linear_attn.in_proj_qkv", binder) + self.assertLess( + binder.index("T > 1 && qkv_quant.has_value()"), + binder.index("slice_la_kkv_blocked(*full_qkv"), + ) + + def test_qwen35_linear_attention_forward_passes_weight_views_to_gemm( + self, + ) -> None: + forward = read_model_file("qwen3_5_forward.cpp") + + for tensor, quant in ( + ("la_in_proj_qkv", "la_in_proj_qkv_quant"), + ("la_in_proj_z", "la_in_proj_z_quant"), + ("la_in_proj_a", "la_in_proj_a_quant"), + ("la_in_proj_b", "la_in_proj_b_quant"), + ("la_out_proj", "la_out_proj_quant"), + ): + self.assertIn( + f"make_weight_view(Lw.{tensor}, Lw.{quant})", + forward, + ) + + +if __name__ == "__main__": + unittest.main()