[Bugfix][RAINFUSION_ATTN] Fix padding-mask rejection on packed [real, pad] layouts and reuse query tail as output padding (#7235)

Signed-off-by: brandneway <gyuan4892@gmail.com>
This commit is contained in:
ooooooye authored and GitHub committed 2026-09-20 10:56:51 +08:00
1 parent fa422a3f44
commit a9abd563f9
4 files changed
+84 -5

No files matched your search

@@ -193,7 +193,11 @@ def test_sparse_quant_uses_public_runtime(runtime, monkeypatch, method):
monkeypatch.setattr(rainfusion_attn, "_mindiesd_supports_precision", lambda: True)
q = torch.randn(1, 4224, 2, 64, dtype=torch.bfloat16)
out = sparse().forward_npu(q, q, q, video_metadata(kv_cache_dtype=method))
assert out.shape == q.shape and not out[:, 4096:].any()
# The sparse path reuses the query tail as the output padding instead of
# zero-filling a fresh buffer; rows past the 4096-row valid prefix are
# never read back downstream (the model unpads by cu_seqlens), so pin the
# passthrough rather than requiring zeros.
assert out.shape == q.shape and torch.equal(out[:, 4096:], q[:, 4096:])
kwargs = runtime.sparse_attention.call_args.kwargs
assert (kwargs["precision"], kwargs["sparse_type"], kwargs["inner_precise"]) == (method, "rf_v3", 4)
runtime.quant_attention.assert_not_called()