diff --git a/tests/models/test_deepseek_v4_vl_rocm.py b/tests/models/test_deepseek_v4_vl_rocm.py index 750927972e..35d5227a6a 100644 --- a/tests/models/test_deepseek_v4_vl_rocm.py +++ b/tests/models/test_deepseek_v4_vl_rocm.py @@ -15,131 +15,6 @@ pytestmark = pytest.mark.skipif( ) -def _preshuffled_fp8_linear( - holder: str = "quant_method", weight_shape: tuple[int, int] = (256, 128) -) -> nn.Module: - pytest.importorskip("aiter") - from vllm.model_executor.kernels.linear.scaled_mm.aiter import ( - AiterPreshuffledFp8BlockScaledMMKernel, - ) - from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import ( - FP8ScaledMMLinearLayerConfig, - ) - from vllm.model_executor.layers.quantization.utils.quant_utils import ( - kFp8Dynamic128Sym, - kFp8Static128BlockSym, - ) - - config = FP8ScaledMMLinearLayerConfig( - weight_quant_key=kFp8Static128BlockSym, - activation_quant_key=kFp8Dynamic128Sym, - input_dtype=torch.bfloat16, - out_dtype=torch.bfloat16, - weight_shape=weight_shape, - ) - layer = nn.Module() - layer.weight = nn.Parameter( - torch.randn(*weight_shape, device="cuda").to(current_platform.fp8_dtype()), - requires_grad=False, - ) - layer.register_parameter( - "weight_scale" if holder == "scheme" else "weight_scale_inv", - nn.Parameter( - torch.ones(*(dim // 128 for dim in weight_shape), device="cuda"), - requires_grad=False, - ), - ) - setattr( - layer, - holder, - SimpleNamespace(fp8_linear=AiterPreshuffledFp8BlockScaledMMKernel(config)), - ) - return layer - - -@pytest.mark.parametrize("holder", ["quant_method", "scheme"]) -def test_rocm_wo_a_keeps_row_major_weights(holder: str, default_vllm_config) -> None: - from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16 - - layer = _preshuffled_fp8_linear(holder) - layer.is_bmm = True - original = layer.weight.detach().clone() - kernel = getattr(layer, holder).fp8_linear - kernel.process_weights_after_loading(layer) - - actual = _get_cached_wo_a_bf16(layer, 2, 128, 128) - torch.testing.assert_close(actual, original.to(torch.bfloat16).view(2, 128, 128)) - - -@pytest.mark.parametrize("holder", ["quant_method", "scheme"]) -def test_rocm_gateup_shuffles_once_and_down_proj_keeps_preshuffled_backend( - holder: str, monkeypatch: pytest.MonkeyPatch, default_vllm_config -) -> None: - from vllm._aiter_ops import rocm_aiter_ops - from vllm.models.deepseek_v4.amd import model as rocm_model - - gate_up, down = _preshuffled_fp8_linear(holder), _preshuffled_fp8_linear(holder) - original_gate_up = gate_up.weight.detach().clone() - original_down = down.weight.detach().clone() - monkeypatch.setattr( - rocm_model, "MergedColumnParallelLinear", lambda *a, **k: gate_up - ) - monkeypatch.setattr(rocm_model, "RowParallelLinear", lambda *a, **k: down) - monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True) - model = rocm_model.DeepseekV4MLP(128, 128, "silu") - for linear in (gate_up, down): - getattr(linear, holder).fp8_linear.process_weights_after_loading(linear) - model.prepare_gateup_preshuffle() - - for linear, original in ((gate_up, original_gate_up), (down, original_down)): - expected = rocm_aiter_ops.shuffle_weight(original, layout=(16, 16)) - torch.testing.assert_close( - linear.weight.view(torch.uint8), expected.view(torch.uint8) - ) - - -def test_rocm_fused_qkv_quant_matches_preshuffled_gemm_scale_layout( - monkeypatch: pytest.MonkeyPatch, default_vllm_config -) -> None: - from vllm._aiter_ops import rocm_aiter_ops - from vllm.models.deepseek_v4.amd.rocm import ( - DeepseekV4ROCMAiterMLAAttention, - apply_pre_quantized_block_scaled_mm, - ) - - if not rocm_aiter_ops.is_blockscale_bpreshuffle_tuned(1024, 256): - pytest.skip("Requires a tuned AITER preshuffled FP8 GEMM") - linear = _preshuffled_fp8_linear(weight_shape=(1024, 256)) - original = linear.weight.detach().clone() - linear.quant_method.fp8_linear.process_weights_after_loading(linear) - attention = DeepseekV4ROCMAiterMLAAttention.__new__(DeepseekV4ROCMAiterMLAAttention) - nn.Module.__init__(attention) - attention.wq_b = linear - attention.indexer = None - attention.q_lora_rank, attention.head_dim = 256, 128 - attention.eps = 1e-5 - attention.q_norm = SimpleNamespace( - weight=torch.ones(256, dtype=torch.bfloat16, device="cuda") - ) - attention.kv_norm = SimpleNamespace( - weight=torch.ones(128, dtype=torch.bfloat16, device="cuda") - ) - monkeypatch.setattr(rocm_aiter_ops, "is_linear_fp8_enabled", lambda: True) - # Multiple rows and groups with different ranges expose byte-order errors. - inputs = torch.randn(16, 384, dtype=torch.bfloat16, device="cuda") - inputs[:, :128] *= 4 - quantized, scales, _ = attention._split_qkv_and_norm(inputs) - assert scales is not None - output = apply_pre_quantized_block_scaled_mm(linear, quantized, scales) - q = inputs[:, :256].float() - normalized = q * torch.rsqrt(q.square().mean(dim=-1, keepdim=True) + attention.eps) - reference_scales = normalized.view(16, 2, 128).abs().amax(dim=-1) - reference_scales /= torch.finfo(current_platform.fp8_dtype()).max - dequantized = quantized.float() * reference_scales.repeat_interleave(128, dim=1) - expected = dequantized @ original.float().T - torch.testing.assert_close(output.float(), expected, atol=0.125, rtol=0.01) - - def test_rocm_packed_kv_cache_auto_uses_ds_mla_layout() -> None: from vllm.config import CacheConfig from vllm.models.deepseek_v4.attention import _resolve_dsv4_kv_cache_dtype diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index da4bc310d7..101e1dd10a 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -175,26 +175,18 @@ def is_aiter_found_and_supported_on_rdna4() -> bool: @functools.cache def _load_gemm_tuned_configs( - csv_path: str, - filters: tuple[tuple[str, object], ...], - key_cols: tuple[str, ...] = ("N", "K", "M"), -) -> set[tuple[int, ...]]: + q_dtype_w: torch.dtype, csv_path: str +) -> set[tuple[int, int, int]]: try: df = pd.read_csv(csv_path).drop_duplicates() - for col, val in filters: - if col not in df.columns: - continue - if isinstance(val, int): - df = df[df[col].astype(int) == val] - else: - df = df[df[col].astype(str) == str(val)] - return set(zip(*(df[c].astype(int) for c in key_cols))) + df = df[df["q_dtype_w"] == str(q_dtype_w)] + return set(zip(df["N"].astype(int), df["K"].astype(int), df["M"].astype(int))) except Exception: return set() def _check_kernel_tuned(N: int, K: int, q_dtype_w: torch.dtype, csv_path: str) -> bool: - configs = _load_gemm_tuned_configs(csv_path, (("q_dtype_w", q_dtype_w),)) + configs = _load_gemm_tuned_configs(q_dtype_w, csv_path) l_m = ( [1, 2, 4] + list(range(8, 513, 8)) @@ -3329,20 +3321,6 @@ class rocm_aiter_ops: csv_path = aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_FILE return _check_kernel_tuned(N, K, q_dtype_w, csv_path) - @staticmethod - def is_blockscale_bpreshuffle_tuned(n: int, k: int) -> bool: - """Whether (N, K) has a tuned aiter blockscale bpreshuffle config.""" - if not current_platform.is_rocm(): - return False - import aiter.ops.gemm_op_a8w8 as aiter_gemm_a8w8_ops - - csv_path = aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_BPRESHUFFLE_FILE - gfx = aiter_gemm_a8w8_ops.get_gfx() - cu_num = aiter_gemm_a8w8_ops.get_cu_num() - return (n, k) in _load_gemm_tuned_configs( - csv_path, (("gfx", gfx), ("cu_num", cu_num)), key_cols=("N", "K") - ) - @staticmethod def shuffle_weight( tensor: torch.Tensor, layout: tuple[int, int] = (16, 16) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 26ad557c67..3c9c616d0c 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -177,7 +177,6 @@ from vllm.model_executor.kernels.linear.scaled_mm.aiter import ( AiterHipbMMPerTokenFp8ScaledMMLinearKernel, AiterInt8ScaledMMLinearKernel, AiterPerTokenFp8ScaledMMLinearKernel, - AiterPreshuffledFp8BlockScaledMMKernel, AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.b12x import ( @@ -322,7 +321,6 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = { "aiter": { AiterInt8ScaledMMLinearKernel, AiterFp8BlockScaledMMKernel, - AiterPreshuffledFp8BlockScaledMMKernel, AiterPerTokenFp8ScaledMMLinearKernel, AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, AiterMxfp4LinearKernel, @@ -466,7 +464,6 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[ BlockWiseTorchFP8ScaledMMLinearKernel, ], PlatformEnum.ROCM: [ - AiterPreshuffledFp8BlockScaledMMKernel, AiterFp8BlockScaledMMKernel, TritonFp8BlockScaledMMKernel, ], @@ -1250,7 +1247,6 @@ __all__ = [ "ScaledMMLinearLayerConfig", "AiterHipbMMPerTokenFp8ScaledMMLinearKernel", "AiterPreshuffledPerTokenFp8ScaledMMLinearKernel", - "AiterPreshuffledFp8BlockScaledMMKernel", "AiterPerTokenFp8ScaledMMLinearKernel", "NvFp4LinearKernel", "NvFp4LinearLayerConfig", diff --git a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py index a57e7510de..03e9be51cf 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py @@ -2,8 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import ClassVar - import torch from vllm import _custom_ops as ops @@ -474,171 +472,3 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): return gemm_a8w8_blockscale_op( A, B, As, Bs, list(self.weight_group_shape), output_dtype=out_dtype ) - - -class AiterPreshuffledFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): - """Aiter FP8 block-scaled GEMM using a pre-shuffled (bpreshuffle) weight.""" - - # gemm_a8w8_blockscale_bpreshuffle reads the activation scale column-major. - wants_transposed_act_scale: ClassVar[bool] = True - - # process_weights_after_loading shuffles layer.weight to layout (16, 16). - preshuffles_weight: ClassVar[bool] = True - - @classmethod - def is_supported( - cls, compute_capability: int | None = None - ) -> tuple[bool, str | None]: - return AiterPreshuffledPerTokenFp8ScaledMMLinearKernel.is_supported( - compute_capability - ) - - @classmethod - def can_implement(cls, config: FP8ScaledMMLinearLayerConfig): - can_implement_base, reason = super().can_implement(config) - if not can_implement_base: - return can_implement_base, reason - - act_quant_desc = config.activation_quant_key.scale - if act_quant_desc.group_shape != GroupShape(1, 128): - return ( - False, - ( - "Supports only dynamic per token group activation " - "quantization with group_shape=(1,128)." - ), - ) - - # bpreshuffle GEMM requires aiter fp8 linear to be enabled and an fp8 - # 2D weight with N and K divisible by 128. Require a tuned config as the - # fallback can trigger faults or numerical errors. - if not rocm_aiter_ops.is_linear_fp8_enabled(): - return ( - False, - ( - "requires setting `VLLM_ROCM_USE_AITER=1` " - "and `VLLM_ROCM_USE_AITER_LINEAR=1`." - ), - ) - - n, k = config.weight_shape - if not (n % 128 == 0 and k % 128 == 0): - return ( - False, - ( - f"requires N and K dimensions divisible by 128, received " - f"N={n} and K={k}." - ), - ) - - if not rocm_aiter_ops.is_blockscale_bpreshuffle_tuned(n, k): - return ( - False, - ( - f"requires a tuned aiter blockscale bpreshuffle config for " - f"N={n} and K={k}." - ), - ) - - return True, None - - @staticmethod - def _reads_weight_directly(layer: torch.nn.Module) -> bool: - """True when something other than apply_weights consumes layer.weight. - - Such a weight must stay in the plain layout: ``is_bmm`` marks a stack - of matrices (wo_a), ``skip_weight_relayout`` marks MLA's kv_b_proj. - Both are stamped after construction, so can_implement cannot see them. - """ - return bool( - getattr(layer, "is_bmm", False) - or getattr(layer, "skip_weight_relayout", False) - ) - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - super().process_weights_after_loading(layer) - - if self._reads_weight_directly(layer): - return - - params = FP8BlockParams.from_layer(layer) - if params.weight_scale_inv is not None: - ws, attr = params.weight_scale_inv, params.WEIGHT_SCALE_INV - else: - ws, attr = params.weight_scale, params.WEIGHT_SCALE - if ws is not None and ws.dtype == torch.float8_e8m0fnu: - replace_parameter(layer, attr, _upcast_e8m0_to_fp32(ws).contiguous()) - - weight = params.weight - # runtime safety net - assert ( - weight.dim() == 2 - and weight.dtype == current_platform.fp8_dtype() - and weight.shape[0] % 128 == 0 - and weight.shape[1] % 128 == 0 - ), ( - "AiterPreshuffledFp8BlockScaledMMKernel requires a 2D fp8 weight " - "with N and K divisible by 128." - ) - - shuffled_weight = rocm_aiter_ops.shuffle_weight( - weight.contiguous(), layout=(16, 16) - ) - replace_parameter( - layer, - params.WEIGHT, - torch.nn.Parameter(shuffled_weight.data, requires_grad=False), - ) - - def apply_weights( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, - **kwargs, - ) -> torch.Tensor: - params = self._get_layer_params(layer) - Bs = ( - params.weight_scale - if params.weight_scale_inv is None - else params.weight_scale_inv - ) - - # Left unshuffled above; kv_b_proj still reaches here from MLA's - # prefill-context path. - plain = self._reads_weight_directly(layer) - - x_2d = x.view(-1, x.shape[-1]) - A, As = rocm_aiter_ops.group_fp8_quant(x_2d, transpose_scale=not plain) - if plain: - output = rocm_aiter_ops.gemm_a8w8_blockscale( - A, - params.weight, - As, - Bs, - list(self.weight_group_shape), - output_dtype=self.config.out_dtype, - ) - else: - output = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle( - A, params.weight, As, Bs, output_dtype=self.config.out_dtype - ) - if bias is not None: - output = output + bias - return output.view(*x.shape[:-1], params.weight.shape[0]) - - def apply_block_scaled_mm( - self, - A: torch.Tensor, - B: torch.Tensor, - As: torch.Tensor, - Bs: torch.Tensor, - ) -> torch.Tensor: - """Block-scaled GEMM for callers that pre-quantize their activations. - - ``As`` must be column-major; a row-major one will not raise, it just - returns wrong numbers for M > 1. - """ - return rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle( - A, B, As, Bs, output_dtype=self.config.out_dtype - ) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 4cbc380bb1..26547c7353 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -582,9 +582,10 @@ class MLAAttention(nn.Module, AttentionLayerBase): ) self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None) self.is_amx_bmm_enabled = getattr(self.impl, "uses_amx_bmm", False) - # MLA reads this weight directly to build W_UK/W_UV, so a backend must - # not relayout it at load time. - kv_b_proj.skip_weight_relayout = True + # AMX reads kv_b_proj's weight directly and never calls it live; the + # reference CPU MLA backend calls it but isn't perf-critical. Skip + # the packed-kernel dispatch either way. + kv_b_proj._cpu_skip_gemm_dispatch = True self.use_direct_call = not current_platform.opaque_attention_op() vllm_config = get_current_vllm_config() diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 9d2e31cfe1..65cda325a2 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -205,9 +205,9 @@ class UnquantizedLinearMethod(LinearMethodBase): def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if current_platform.is_cpu(): - # MLA's kv_b_proj (see `skip_weight_relayout`): not perf-critical, - # so skip packing and use a plain fallback. - if getattr(layer, "skip_weight_relayout", False): + # MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): not + # perf-critical, so skip packing and use a plain fallback. + if getattr(layer, "_cpu_skip_gemm_dispatch", False): layer.cpu_linear = torch.nn.functional.linear return diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index c1a6903937..b23e728b46 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -71,10 +71,7 @@ from vllm.model_executor.models.utils import ( make_layers, maybe_prefix, ) -from vllm.models.deepseek_v4.amd.rocm import ( - DeepseekV4ROCMAiterMLAAttention, - weight_already_preshuffled, -) +from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention from vllm.platforms import current_platform from vllm.platforms.rocm import on_gfx950 from vllm.sequence import IntermediateTensors @@ -152,13 +149,11 @@ class DeepseekV4MLP(nn.Module): return if ws.dtype == torch.float8_e8m0fnu: ws = _upcast_e8m0_to_fp32(ws).contiguous() - # Skip if the linear's kernel already shuffled it. - if not weight_already_preshuffled(self.gate_up_proj): - replace_parameter( - self.gate_up_proj, - "weight", - rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), - ) + replace_parameter( + self.gate_up_proj, + "weight", + rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), + ) self._gateup_scale = ws def forward(self, x): diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 3cc9827c1f..f8ac39cd62 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -62,21 +62,6 @@ def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor: return indptr -def weight_already_preshuffled(linear: torch.nn.Module) -> bool: - """True when the linear's kernel already B-preshuffled ``weight``. - - The hand-shuffles below (fused_wqa_wkv, wo_b, gate_up_proj) must be skipped - for those, since shuffle_weight is a permutation rather than an involution. - """ - return any( - getattr(getattr(method, "fp8_linear", None), "preshuffles_weight", False) - for method in ( - getattr(linear, "quant_method", None), - getattr(linear, "scheme", None), - ) - ) - - def apply_pre_quantized_block_scaled_mm( linear: torch.nn.Module, x_fp8: torch.Tensor, @@ -792,13 +777,12 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): return None if ws.dtype == torch.float8_e8m0fnu: ws = _upcast_e8m0_to_fp32(ws).contiguous() - # Skip if the linear's kernel already shuffled it. - if not weight_already_preshuffled(linear): - replace_parameter( - linear, - "weight", - rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), - ) + # Shuffle the weight in place (single weight, no unshuffled copy). + replace_parameter( + linear, + "weight", + rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), + ) return ws self._wqa_wkv_scale = _prep(self.fused_wqa_wkv) @@ -896,26 +880,16 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): return qr_kv, kv_score, indexer_kv_score, indexer_weights @functools.cached_property - def _wq_b_act_scale_transpose(self) -> bool | None: - """Activation-scale byte order every wq_b consumer agrees on, else None. - - None means "do not take the fused norm+quant path". Otherwise the value - is the ``transpose_scale`` the producer must pass so that the scale it - writes matches what the selected GEMM reads. + def _wq_b_uses_aiter_block_scaled(self) -> bool: + """True when both wq_b GEMMs run the aiter block-scaled fp8 kernel. Cached: the linear kernels and the aiter env gates are fixed once the model is built, so this is evaluated at the first forward only. - Two conditions, both necessary: - - * The fused norm+quant path is only valid if the quant and GEMM it - replaces are exactly the aiter ones; otherwise fall back to the - shared path. - * One qr_scale feeds both self.wq_b and self.indexer.wq_b, which have - different (N, K) and so can resolve to kernels wanting opposite byte - orders. A single producer cannot serve both, so refuse the fast path - when they disagree rather than guessing. + The fused norm+quant path is only valid if the quant and GEMM it + replaces are exactly the aiter ones; otherwise fall back to the + shared path. """ from vllm._aiter_ops import rocm_aiter_ops from vllm.model_executor.kernels.linear.scaled_mm import ( @@ -923,31 +897,16 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): ) if not rocm_aiter_ops.is_linear_fp8_enabled(): - return None + return False linears = [self.wq_b] if self.indexer is not None: linears.append(self.indexer.wq_b) - layouts: set[bool] = set() for linear in linears: kernel = getattr(getattr(linear, "quant_method", None), "fp8_linear", None) if not isinstance(kernel, Fp8BlockScaledMMLinearKernel): - return None - layouts.add(bool(getattr(kernel, "wants_transposed_act_scale", False))) - if len(layouts) != 1: - logger.warning_once( - "DeepSeek-V4 wq_b consumers disagree on activation-scale layout; " - "disabling the fused q/kv norm+quant path.", - scope="global", - ) - return None - transpose_scale = layouts.pop() - logger.debug_once( - "DeepSeek-V4 wq_b: emitting %s-major activation scales", - "column" if transpose_scale else "row", - scope="global", - ) - return transpose_scale + return False + return True def _split_qkv_and_norm( self, qr_kv: torch.Tensor @@ -964,12 +923,11 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): when the aiter linear path is not active. """ qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1) - transpose_scale = self._wq_b_act_scale_transpose if not ( qr.dim() == 2 and qr.shape[0] > 0 and self.q_lora_rank % 128 == 0 - and transpose_scale is not None + and self._wq_b_uses_aiter_block_scaled ): return super()._split_qkv_and_norm(qr_kv) @@ -983,8 +941,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): kv_weight=self.kv_norm.weight.data, kv_epsilon=self.eps, group_size=128, - # Emit the byte order the wq_b GEMMs read. - transpose_scale=transpose_scale, + transpose_scale=False, ) def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: