[triton][mqa] fix silent tail-row drop in deepgemm_fp8_paged_mqa_logits at large output stride

[zejunchen-zejun]

At large MTP batches the sparse-indexer decode path issues one big deepgemm_fp8_paged_mqa_logits + top_k_per_row_decode launch whose later rows can retain stale sparse-KV indices in the persistent buffer (rows beyond the CU wave capacity), so requests past the cliff attend to wrong KV and their drafts get rejected -> acceptance halves (~50% -> ~25% at con256, cliff at seq 128), so when specify max num seqs to the value which is smaller than 128, the accept ratio restores to the normal value, while for default 512, most of requests has low draft token accept ratio. Here is the histogram illustration.

Root cause: the Gluon kernel stores logits with gl.amd.cdna3.buffer_store, whose AMD hardware voffset is a 32-bit byte offset. The output address for row r is r * out_logits.stride(0) * out_logits.element_size(). Once that reaches 2**31 the offset overflows and the store is silently dropped, leaving the tail rows unwritten (they keep their prior/zero contents).

This bites callers that allocate a wide dense logits tensor. Concretely, GLM-5.2 sparse-MLA (DSA) MTP decode allocates logits as [batchnext_n, max_model_len] with max_model_len = 1<<20; at con>=256 with next_n=4 that is [1024, 1<<20], so stride(0)=1<<20 and row 512 hits exactly 512 * (1<<20) * 4 = 2**31 bytes. Rows 512..1023 are never written -> top-k reads all-zero rows -> wrong sparse-KV indices -> MTP acceptance collapses (~50% -> ~25%). Verified with a standalone reproducer: with physical_cols=1<<20 the first untouched row is 512; with physical_cols=600000 it moves to 895 (= where rstride*4 first crosses 2**31), proving it is a byte-offset boundary, not a fixed row count.

Fix
Address the output in 64-bit by advancing the base pointer, keeping the buffer_store voffset a small int32 column index:

• gluon/pa_mqa_logits.py — all 19 stores across the 3 gluon kernels changed from ptr=OutLogits_buffer, offsets=row*stride_out_batch + col to ptr=OutLogits_buffer + row.to(tl.int64)*stride_out_batch, offsets=col (int32 columns only).
• attention/pa_mqa_logits.py (wrapper) — declare stride_out_batch as i64 and drop the tt.pointer_range 32 hint on OutLogits_buffer so the base may be addressed beyond 2 GB.
• _triton_kernels/attention/pa_mqa_logits.py (non-gluon path) — type all stride_out_batch params tl.int64 so row * stride_out_batch promotes to int64 before tl.store (which already does 64-bit addressing); matches the one kernel there that was already i64.
  This is a single-launch fix — no chunking, no extra kernel launches, no change to the non-overflowing path.

Here is the validation result by model side:

atom

Configuration	Acceptance Rate	Accepted 0 Tokens	Accepted 1 Token	Accepted 2 Tokens	Accepted 3 Tokens
without fix	26.87%	52.30%	23.90%	14.70%	9.10%
with fix	53.53%	17.64%	27.65%	31.20%	23.51%

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	20	exact_match	↑	0.931	±	0.007
		strict-match	20	exact_match	↑	0.931	±	0.007

atom-vllm

Configuration	Position 1 Acceptance Rate	Position 2 Acceptance Rate	Position 3 Acceptance Rate	Avg Draft Acceptance Rate
without fix	0.422	0.245	0.108	25.8%
with fix	0.721	0.494	0.235	48.3%

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	20	exact_match	↑	0.9274	±	0.0076
		strict-match	20	exact_match	↑	0.9289	±	0.0075

small reproducer:

#!/usr/bin/env python3
"""Reproduce the paged-MQA logits row coverage issue without ATOM/vLLM."""

from __future__ import annotations

import argparse

import torch

from aiter import dtypes
from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits


GREEN = "\033[32m"
RED = "\033[31m"
RESET = "\033[0m"


def _green(text: str) -> str:
    return f"{GREEN}{text}{RESET}"


def _red(text: str) -> str:
    return f"{RED}{text}{RESET}"


def _print_scalar(name: str, value: object) -> None:
    print(f"  {name}: {value}", flush=True)


def _print_tensor(name: str, tensor: torch.Tensor) -> None:
    print(
        f"  {name}: shape={tuple(tensor.shape)}, dtype={tensor.dtype}, "
        f"device={tensor.device}, stride={tensor.stride()}, "
        f"is_contiguous={tensor.is_contiguous()}",
        flush=True,
    )


def _make_inputs(batch_size: int, next_n: int, context_len: int, *, device: str):
    heads = 128
    head_dim = 128
    block_size = 256
    max_block_len = (context_len + block_size - 1) // block_size
    num_blocks = max_block_len

    # Use non-zero FP8 bit patterns; exact values are irrelevant for row coverage.
    q_bits = torch.randint(
        1,
        64,
        (batch_size, next_n, heads, head_dim),
        dtype=torch.uint8,
        device=device,
    )
    q_fp8 = q_bits.view(dtypes.fp8)

    kv_bits = torch.randint(
        1,
        64,
        (num_blocks, block_size, 1, head_dim + 4),
        dtype=torch.uint8,
        device=device,
    )
    # The last 4 fp8 bytes per token are viewed as one fp32 scale by the kernel.
    # Force a finite positive scale value: float32(1.0) bit pattern.
    scale_bytes = torch.tensor([0, 0, 128, 63], dtype=torch.uint8, device=device)
    kv_bits[..., head_dim:] = scale_bytes
    kv_cache = kv_bits.view(dtypes.fp8)

    weights = torch.ones((batch_size * next_n, heads), dtype=torch.float32, device=device)
    context_lens = torch.full((batch_size,), context_len, dtype=torch.int32, device=device)
    block_tables = torch.arange(max_block_len, dtype=torch.int32, device=device).repeat(
        batch_size, 1
    )
    return q_fp8, kv_cache, weights, context_lens, block_tables, block_size


def _run_kernel(
    label: str,
    q_fp8: torch.Tensor,
    kv_cache: torch.Tensor,
    weights: torch.Tensor,
    context_lens: torch.Tensor,
    block_tables: torch.Tensor,
    block_size: int,
    max_model_len_arg: int,
):
    device = "cuda"
    batch_size, next_n, _, _ = q_fp8.shape
    rows = batch_size * next_n
    context_len = int(context_lens.max().item())
    sentinel = 12345.0
    out = torch.full(
        (rows, max_model_len_arg), sentinel, dtype=torch.float32, device=device
    )
    print(f"\n[{label}] deepgemm_fp8_paged_mqa_logits inputs", flush=True)
    _print_scalar("batch_size", batch_size)
    _print_scalar("next_n", next_n)
    _print_scalar("rows", rows)
    _print_scalar("context_len", context_len)
    _print_scalar("max_model_len_arg", max_model_len_arg)
    _print_scalar("KVBlockSize", block_size)
    _print_scalar("ChunkK", 256)
    _print_scalar("WavePerEU", 2)
    _print_scalar("Preshuffle", True)
    _print_tensor("q_fp8", q_fp8)
    _print_tensor("kv_cache", kv_cache)
    _print_tensor("weights", weights)
    _print_tensor("out_logits", out)
    _print_tensor("context_lens", context_lens)
    _print_tensor("block_tables", block_tables)

    deepgemm_fp8_paged_mqa_logits(
        q_fp8,
        kv_cache,
        weights,
        out,
        context_lens,
        block_tables,
        max_model_len_arg,
        Preshuffle=True,
        KVBlockSize=block_size,
        ChunkK=256,
        WavePerEU=2,
    )
    torch.cuda.synchronize()
    row_touched = (out != sentinel).any(dim=1)
    touched_rows = int(row_touched.sum().item())
    first_untouched = (
        int((~row_touched).nonzero()[0].item())
        if not bool(row_touched.all().item())
        else None
    )
    print(
        f"[{label}] result: touched_rows={touched_rows}, "
        f"first_untouched_row={first_untouched}",
        flush=True,
    )
    return out, row_touched


def _print_touch_check(label: str, row_touched: torch.Tensor, expected_rows: int) -> None:
    touched_rows = int(row_touched.sum().item())
    first_untouched = (
        int((~row_touched).nonzero()[0].item())
        if not bool(row_touched.all().item())
        else None
    )
    if touched_rows == expected_rows:
        print(
            _green(
                f"OK {label}: touched_rows={touched_rows}/{expected_rows}, "
                f"first_untouched_row={first_untouched}"
            ),
            flush=True,
        )
    else:
        print(
            _red(
                f"BAD {label}: touched_rows={touched_rows}/{expected_rows}, "
                f"first_untouched_row={first_untouched}"
            ),
            flush=True,
        )


def _print_compare(label: str, actual: torch.Tensor, expected: torch.Tensor) -> None:
    equal = bool(torch.equal(actual, expected))
    if equal:
        print(_green(f"OK {label}: exact match with compact reference"), flush=True)
        return

    mismatches = int((actual != expected).sum().item())
    finite = torch.isfinite(actual) & torch.isfinite(expected)
    max_abs_diff = None
    if bool(finite.any().item()):
        max_abs_diff = float((actual[finite] - expected[finite]).abs().max().item())
    print(
        _red(
            f"BAD {label}: mismatch with compact reference, "
            f"mismatched_elements={mismatches}, max_abs_diff={max_abs_diff}"
        ),
        flush=True,
    )


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--context-len", type=int, default=112592)
    parser.add_argument("--physical-cols", type=int, default=1048576)
    args = parser.parse_args()

    torch.manual_seed(1234)
    next_n = 4
    context_len = args.context_len
    physical_cols = args.physical_cols
    if physical_cols < context_len:
        raise ValueError("--physical-cols must be >= --context-len")

    q_fp8, kv_cache, weights, context_lens, block_tables, block_size = _make_inputs(
        256, next_n, context_len, device="cuda"
    )

    reference, reference_touched = _run_kernel(
        "reference_compact_batch256_rows1024",
        q_fp8,
        kv_cache,
        weights,
        context_lens,
        block_tables,
        block_size,
        context_len,
    )
    one_shot, one_shot_touched = _run_kernel(
        "one_shot_batch256_rows1024",
        q_fp8,
        kv_cache,
        weights,
        context_lens,
        block_tables,
        block_size,
        physical_cols,
    )
    chunk0, chunk0_touched = _run_kernel(
        "chunk0_batch128_rows512",
        q_fp8[:128],
        kv_cache,
        weights[:512],
        context_lens[:128],
        block_tables[:128],
        block_size,
        physical_cols,
    )
    chunk1, chunk1_touched = _run_kernel(
        "chunk1_batch128_rows512",
        q_fp8[128:],
        kv_cache,
        weights[512:],
        context_lens[128:],
        block_tables[128:],
        block_size,
        physical_cols,
    )

    one_shot_first_untouched = (
        int((~one_shot_touched).nonzero()[0].item())
        if not bool(one_shot_touched.all().item())
        else None
    )
    untouched_rows = [int(x) for x in (~one_shot_touched).nonzero().flatten()[:16].tolist()]

    print("\nTouch Checks", flush=True)
    _print_touch_check("reference compact", reference_touched, 1024)
    _print_touch_check("one-shot large-stride", one_shot_touched, 1024)
    _print_touch_check("chunk0 large-stride", chunk0_touched, 512)
    _print_touch_check("chunk1 large-stride", chunk1_touched, 512)

    print("\nAccuracy Checks", flush=True)
    _print_compare(
        "one-shot rows 0..511 valid columns",
        one_shot[:512, :context_len],
        reference[:512],
    )
    _print_compare(
        "one-shot rows 512..1023 valid columns",
        one_shot[512:, :context_len],
        reference[512:],
    )
    _print_compare(
        "chunk0 rows 0..511 valid columns",
        chunk0[:, :context_len],
        reference[:512],
    )
    _print_compare(
        "chunk1 rows 512..1023 valid columns",
        chunk1[:, :context_len],
        reference[512:],
    )

    print("\nSummary", flush=True)
    _print_scalar("case", "deepgemm_fp8_paged_mqa_logits_row_coverage")
    _print_scalar("one_shot_rows", 1024)
    _print_scalar("one_shot_context_len", context_len)
    _print_scalar("one_shot_physical_cols", physical_cols)
    _print_scalar("one_shot_touched_rows", int(one_shot_touched.sum().item()))
    _print_scalar("one_shot_first_untouched_row", one_shot_first_untouched)
    _print_scalar("one_shot_untouched_rows_first_16", untouched_rows)
    _print_scalar("chunk0_touched_rows", int(chunk0_touched.sum().item()))
    _print_scalar("chunk1_touched_rows", int(chunk1_touched.sum().item()))
    _print_scalar("chunk0_all_touched", bool(chunk0_touched.all().item()))
    _print_scalar("chunk1_all_touched", bool(chunk1_touched.all().item()))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4244 --add-label <label>

[Copilot]

Pull request overview

This PR fixes a correctness bug in the paged-MQA FP8 logits Triton path where AMD buffer_store silently drops stores once the per-row output byte offset crosses the 2**31 boundary, leaving tail rows unwritten for very wide dense logits tensors.

Changes:

• Update Gluon buffer_store addressing to advance the base pointer in 64-bit (ptr += row*stride) while keeping voffset small (column-only), preventing 32-bit byte-offset overflow drops.
• Widen stride_out_batch to 64-bit in the wrapper signature and ensure non-Gluon kernels treat stride_out_batch as tl.int64.
• Add a regression test that crosses the 2**31 byte-offset boundary and asserts both “all rows touched” and bit-identical results vs a compact reference.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
op_tests/test_pa_mqa_logits_offset.py	Adds a regression test that reproduces the tail-row drop at wide output stride and validates row coverage + correctness.
aiter/ops/triton/gluon/pa_mqa_logits.py	Fixes Gluon buffer_store writes by using 64-bit base-pointer advancement per output row.
aiter/ops/triton/attention/pa_mqa_logits.py	Updates kernel signature to use i64 for stride_out_batch and removes tt.pointer_range constraint for the output pointer.
aiter/ops/triton/_triton_kernels/attention/pa_mqa_logits.py	Types stride_out_batch as tl.int64 so row-stride address arithmetic promotes to 64-bit.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[valarLip]

LGTM

[k50112113]

LGTM