Skip to content

[GG] Revert unconditional MLA BMM contiguity workaround #138

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
voipmonitor merged 1 commit into from
Jul 20, 2026
Merged

[GG] Revert unconditional MLA BMM contiguity workaround #138

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
131 changes: 0 additions & 131 deletions tests/v1/attention/test_mla_backends.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,77 +86,6 @@ def test_mla_kv_cache_spec_uses_layer_cache_dtype(
assert spec.page_size_bytes == 64 * 656


def test_mla_init_propagates_backend_bmm_contiguity_contract(monkeypatch):
class FakeImpl:
is_sparse = True
supports_mha_prefill = False

def __init__(self, **kwargs):
self.force_contiguous_mla_bmm_input = True
self.force_contiguous_mla_bmm_weight = True
self.force_contiguous_mla_bmm_output = True

class FakeBackend:
@staticmethod
def is_mla():
return True

@staticmethod
def get_name():
return "TEST_MLA"

@staticmethod
def get_impl_cls():
return FakeImpl

config = SimpleNamespace(
compilation_config=SimpleNamespace(
static_forward_context={}, cudagraph_capture_sizes=[]
),
attention_config=SimpleNamespace(mla_prefill_backend=None),
parallel_config=SimpleNamespace(decode_context_parallel_size=1),
scheduler_config=SimpleNamespace(max_num_batched_tokens=128),
)
monkeypatch.setattr(
mla_attention_module, "get_current_vllm_config", lambda: config
)
monkeypatch.setattr(
mla_attention_module, "get_current_vllm_config_or_none", lambda: config
)
monkeypatch.setattr(mla_attention_module, "_init_kv_cache_quant", lambda *a: None)
monkeypatch.setattr(
mla_attention_module,
"get_mla_prefill_backend",
lambda _: (_ for _ in ()).throw(ValueError),
)
monkeypatch.setattr(mla_attention_module, "_DecodeConcatQuantFP8", lambda **_: None)
monkeypatch.setattr(mla_attention_module, "QuantFP8", lambda **_: None)
monkeypatch.setattr(
mla_attention_module.rocm_aiter_ops, "is_fp8bmm_enabled", lambda: False
)
monkeypatch.setattr(
mla_attention_module.rocm_aiter_ops, "is_fp4bmm_enabled", lambda: False
)

layer = MLAAttention(
num_heads=8,
scale=1.0,
qk_nope_head_dim=4,
qk_rope_head_dim=2,
v_head_dim=3,
q_lora_rank=None,
kv_lora_rank=4,
kv_b_proj=SimpleNamespace(),
prefix="test",
attn_backend=FakeBackend,
use_sparse=True,
)

assert layer.force_contiguous_mla_bmm_input
assert layer.force_contiguous_mla_bmm_weight
assert layer.force_contiguous_mla_bmm_output


# Remove sm100 backends from the list if not using sm100
if not torch.cuda.is_available() or torch.cuda.get_device_properties(0).major < 10:
BACKENDS_TO_TEST.remove(AttentionBackendEnum.CUTLASS_MLA)
Expand Down Expand Up @@ -193,7 +122,6 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch):
layer.kv_b_proj.quant_method = None
layer.is_aiter_triton_fp4_bmm_enabled = False
layer.is_aiter_triton_fp8_bmm_enabled = False
layer.force_contiguous_mla_bmm_weight = False
layer.quant_config = None
layer.layer_name = "test"

Expand All @@ -219,65 +147,6 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch):
torch.testing.assert_close(layer.W_UK_T, old_w_uk_t + 100)


def test_mla_post_load_honors_bmm_weight_contiguity(monkeypatch):
layer = MLAAttention.__new__(MLAAttention)
torch.nn.Module.__init__(layer)
layer.kv_lora_rank = 2
layer.num_heads = 2
layer.qk_nope_head_dim = 3
layer.v_head_dim = 4
layer.kv_b_proj = torch.nn.Module()
layer.kv_b_proj.weight = torch.nn.Parameter(
torch.arange(28.0, dtype=torch.float32).reshape(14, 2)
)
layer.kv_b_proj.quant_method = None
layer.is_aiter_triton_fp4_bmm_enabled = False
layer.is_aiter_triton_fp8_bmm_enabled = False
layer.force_contiguous_mla_bmm_weight = True
layer.quant_config = None
layer.layer_name = "test"

monkeypatch.setattr(
mla_attention_module, "set_default_quant_scales", lambda *_, **__: None
)

with torch.no_grad():
layer.process_weights_after_loading(torch.float32)

assert layer.W_UV.is_contiguous()
assert layer.W_UK_T.is_contiguous()


def test_mla_v_up_proj_honors_bmm_contiguity(monkeypatch):
layer = object.__new__(MLAAttention)
layer.num_heads = 8
layer.kv_lora_rank = 4
layer.v_head_dim = 3
layer.is_aiter_triton_fp4_bmm_enabled = False
layer.is_aiter_triton_fp8_bmm_enabled = False
layer.force_contiguous_mla_bmm_input = True
layer.force_contiguous_mla_bmm_output = True
layer.W_UV = torch.randn((8, 4, 3), dtype=torch.bfloat16)

x = torch.randn((6, 8, 4), dtype=torch.bfloat16)
out = torch.empty((6, 8, 3), dtype=torch.bfloat16)
expected = torch.einsum("bnl,nlv->bnv", x, layer.W_UV)
real_bmm = torch.bmm
seen_layouts = []

def checked_bmm(input_tensor, mat2, *, out=None):
assert out is not None
seen_layouts.append((input_tensor.is_contiguous(), out.is_contiguous()))
return real_bmm(input_tensor, mat2, out=out)

monkeypatch.setattr(torch, "bmm", checked_bmm)

MLAAttention._v_up_proj(layer, x, out)

assert seen_layouts == [(True, True)]
torch.testing.assert_close(out, expected)


# Filtered per-test via validate_configuration (capability/deps/dims).
PREFILL_BACKENDS_TO_TEST = [
MLAPrefillBackendEnum.FLASH_ATTN,
Expand Down
37 changes: 3 additions & 34 deletions vllm/model_executor/layers/attention/mla_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -604,15 +604,6 @@ def __init__(
**extra_impl_args,
)
self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None)
self.force_contiguous_mla_bmm_input = getattr(
self.impl, "force_contiguous_mla_bmm_input", False
)
self.force_contiguous_mla_bmm_weight = getattr(
self.impl, "force_contiguous_mla_bmm_weight", False
)
self.force_contiguous_mla_bmm_output = getattr(
self.impl, "force_contiguous_mla_bmm_output", False
)
self.use_direct_call = not current_platform.opaque_attention_op()

vllm_config = get_current_vllm_config()
Expand Down Expand Up @@ -1098,9 +1089,6 @@ def forward_impl(
N, B, P = mqa_q_nope.shape
_, _, L = self.W_UK_T.shape

if self.force_contiguous_mla_bmm_input:
mqa_q_nope = mqa_q_nope.contiguous()

if self.q_pad_num_heads is not None:
mqa_ql_nope = mqa_q_nope.new_empty((self.q_pad_num_heads, B, L))
mqa_ql_nope.resize_((N, B, L))
Expand Down Expand Up @@ -1500,14 +1488,9 @@ def process_weights_after_loading(self, act_dtype: torch.dtype):
)
else:
# Convert from (L, N, V) to (N, L, V)
W_UV = W_UV.transpose(0, 1)
replace_parameter(self, "W_UV", W_UV.transpose(0, 1), prefer_copy=True)
# Convert from (L, N, P) to (N, P, L)
W_UK_T = W_UK.permute(1, 2, 0)
if self.force_contiguous_mla_bmm_weight:
W_UV = W_UV.contiguous()
W_UK_T = W_UK_T.contiguous()
replace_parameter(self, "W_UV", W_UV, prefer_copy=True)
replace_parameter(self, "W_UK_T", W_UK_T, prefer_copy=True)
replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True)

# If we should not load quant weights, we initialize the scales to 1.0
# as the default value. See [Note: Register q/k/v/prob scales in state dict]
Expand Down Expand Up @@ -1616,21 +1599,7 @@ def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor):
)
else:
# Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V)
# Some CUDA BMM algorithms read a full tile beyond strided tensor
# bounds. Backends with tightly mapped buffers opt into contiguous
# operands so those accesses stay inside the logical allocation.
if self.force_contiguous_mla_bmm_input:
x = x.contiguous()
if self.force_contiguous_mla_bmm_output:
bmm_out = torch.empty(
(self.num_heads, x.shape[1], self.v_head_dim),
dtype=out.dtype,
device=out.device,
)
torch.bmm(x, self.W_UV, out=bmm_out)
out.copy_(bmm_out.transpose(0, 1))
else:
torch.bmm(x, self.W_UV, out=out.transpose(0, 1))
torch.bmm(x, self.W_UV, out=out.transpose(0, 1))

def _v_up_proj_bmm(
self,
Expand Down
Loading