Skip to content

Fix fused_mla_rope_kv_split for v_dim == 0 - #6340

Open
htesd wants to merge 1 commit into
NVIDIA:devfrom
htesd:fix/fused-mla-rope-kv-split-vdim0
Open

Fix fused_mla_rope_kv_split for v_dim == 0#6340
htesd wants to merge 1 commit into
NVIDIA:devfrom
htesd:fix/fused-mla-rope-kv-split-vdim0

Conversation

@htesd

@htesd htesd commented Aug 7, 2026

Copy link
Copy Markdown

What does this PR do?

fused_mla_rope_kv_split fails when v_dim == 0. This is a legitimate
configuration for callers that only need the rotary-embedded key stream
from the fused kernel (e.g. compressed-attention paths that carry V
separately).

There are two independent failures:

  1. Wrapper (sbhd): the value output (forward) and the value
    gradient (backward) tensors have zero elements, so
    view(..., -1, ...) cannot infer the remaining dimension and raises
    RuntimeError: cannot reshape tensor of 0 elements .... Fixed by
    using the explicit batch_size in the forward reshape and dk's
    flattened length in the backward reshape. The thd path derives
    shapes from cu_seqlens and is unaffected.

  2. Triton kernels: the value load/store offsets are built with
    tl.arange(0, v_dim) unconditionally, and tl.arange requires
    end > start, so kernel specialization fails at compile time with
    ValueError: arange's end argument must be greater than the start argument. Fixed by guarding the value path behind the v_dim
    constexpr — the branch is resolved at compile time, so for
    v_dim > 0 the generated kernel is unchanged.

Tests

Adds a v_dim == 0 case to the existing kv_split forward/backward
parity tests in tests/unit_tests/fusions/test_mla_yarn_rope_apply.py
(covers both sbhd and thd input formats). The new case fails on current
dev (wrapper raises first; with the wrapper fixed, kernel
specialization still fails) and passes with this fix. The pre-existing
kv_split cases pass unchanged (verified on H200 / sm_90a).

Two independent failures break v_dim == 0 (callers that only need the
rotary-embedded key stream from the fused kernel):

1. Wrapper (sbhd): the value output (forward) and value gradient
   (backward) tensors have zero elements, so view(..., -1, ...) cannot
   infer the remaining dimension and raises. Use the explicit
   batch_size in the forward reshape and dk's flattened length in the
   backward reshape. The thd path derives shapes from cu_seqlens and is
   unaffected.

2. Triton kernels: the value load/store offsets are built with
   tl.arange(0, v_dim) unconditionally, and tl.arange requires
   end > start, so kernel specialization fails at compile time. Guard
   the value path behind the v_dim constexpr; for v_dim > 0 the kernel
   behavior is unchanged (the branch is resolved at compile time).

Adds a v_dim == 0 case to the existing kv_split unit tests (covers
both sbhd and thd input formats).

Signed-off-by: iiap <1471127927@qq.com>
@htesd
htesd requested review from a team as code owners August 7, 2026 08:38
@copy-pr-bot

copy-pr-bot Bot commented Aug 7, 2026

Copy link
Copy Markdown

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants