[HuiyingLi]

What does this PR do ?

Fix DeepSeek-V4 TileLang sparse-MLA backward compilation on Blackwell GPUs.

TileLang 0.1.11 vectorizes the per-block mask fragment. When the fragment is
stored as bool, CUDA code generation attempts to materialize boolx8, which
has no supported CUDA storage type. The Python interface already normalizes the
validity mask to 0/1 int32, so this change preserves that representation in
the kernel and forms scalar predicates only where they are consumed.

Changelog

• Store the sparse-MLA backward validity fragment as int32 instead of bool.
• Sum the normalized 0/1 mask directly when detecting all-invalid blocks.
• Preserve the existing masking semantics when computing attention probabilities.

Validation

• 1x GB200, TileLang 0.1.11, AutoModel 26.06 container:
  python3 -m pytest -q -s tests/unit_tests/models/deepseek_v4/test_dsv4_layers.py -k "test_sparse_attention_tilelang_backend_matches_torch"
  • 2 passed, including the causal-padding shape.
  • Both cases compiled sparse_mqa_bwd_kernel and compared forward output plus Q, KV, and attention-sink gradients against the torch reference.
  • Slurm allocation: 5987874.
• CPU DSV4 layer suite: 70 passed, 10 skipped.
• ruff check, ruff format --check, and git diff --check: passed.

Before your PR is "Ready for review"

Pre checks:

• Read and followed the contributor guidelines.
• Existing forward/backward parity tests cover the changed kernel contract.
• No documentation update is required for this internal code-generation fix.

Additional Information

• The change is isolated to the vendored DeepSeek-V4 TileLang sparse-MLA backward kernel.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[HuiyingLi]

/claude review

[HuiyingLi]

/ok to test 77b0a5c

[claude]

LGTM

Reviewed the single changed file tilelang_sparse_mla_bwd.py (excluding uv.lock; none present). The change swaps the mask fragment from bool to T.int32 to avoid materializing boolx8 (no CUDA storage type on Blackwell under TileLang 0.1.11's vectorized fragment store).

Verified numerical equivalence to the prior behavior:

• ValidMask is produced by the Python interface as exactly 0/1 int32 (((topk_idxs >= 0) & (topk_idxs < S_kv)).to(torch.int32)), so storing it directly into the int32 fragment (dropping the redundant != 0) preserves values.
• valid_count[0] += mask[bi_i] summing 0/1 int32 is identical to the previous if_then_else(mask, 1, 0).
• The scalar boolean predicate is correctly re-formed as mask[bi_i] != 0 at the one if_then_else use site.

No gradient-determinism or dKV-accumulation behavior changed. As a Blackwell-specific codegen workaround with unchanged numerics, the absence of a new CPU unit test is appropriate.

[akoumpa]

/ok to test 0ffbfc2