A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 01:00 UTC
[ROCm] Fix TurboQuant on ROCm: backend routing, flash-attn compat, int64 overflow (#39953)
Signed-off-by: aditi <aditi.rana@amd.com>
aditi-amd committed
Apr 17, 2026 at 13:08 UTC
6ef1efd51f11106fc44deb9e7b2f5cd1247fc37e
4 files changed
+22
-23
vllm/platforms/rocm.py
+1
@@ -382,6 +382,7 @@ def _get_backend_priorities(
382
if is_aiter_found_and_supported():
383
backends.append(AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN)
384
backends.append(AttentionBackendEnum.TRITON_ATTN)
385
+ backends.append(AttentionBackendEnum.TURBOQUANT)
386
387
return backends
388
vllm/v1/attention/backends/turboquant_attn.py
+3
-11
@@ -507,8 +507,7 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
507
# max_query_len == max_seq_len means no request has prior cached KV.
508
# Both are Python ints — no GPU sync.
509
if _HAS_FLASH_ATTN and attn_metadata.max_query_len == attn_metadata.max_seq_len:
510
- output = torch.empty(N, Hq, D, device=query.device, dtype=query.dtype)
511
- flash_attn_varlen_func(
510
+ return flash_attn_varlen_func(
511
q=query,
512
k=key,
513
v=value,
@@ -518,9 +517,7 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
517
max_seqlen_k=attn_metadata.max_query_len,
518
softmax_scale=self.scale,
519
causal=True,
521
- out=output,
520
)
523
- return output
521
522
# Continuation or no flash_attn: per-request attention.
523
# For continuation chunks (seq_len > q_len), we must attend to
@@ -557,10 +554,9 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
554
if q_len == seq_len:
555
# First-chunk prefill: all K/V are in the current batch.
556
if _HAS_FLASH_ATTN:
560
- out = torch.empty_like(q_seq)
557
_cu_2[1] = q_len
558
cu = _cu_2
563
- flash_attn_varlen_func(
559
+ out = flash_attn_varlen_func(
560
q=q_seq,
561
k=k_seq,
562
v=v_seq,
@@ -570,7 +566,6 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
566
max_seqlen_k=q_len,
567
softmax_scale=self.scale,
568
causal=True,
573
- out=out,
569
)
570
else:
571
q_t = q_seq.transpose(0, 1).contiguous()
@@ -733,10 +728,9 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
728
729
# Attention: q_len queries attending to seq_len K/V with causal mask
730
if _HAS_FLASH_ATTN:
736
- output = torch.empty(q_len, Hq, D, device=device, dtype=query.dtype)
731
cu_seqlens_q = torch.tensor([0, q_len], device=device, dtype=torch.int32)
732
cu_seqlens_k = torch.tensor([0, seq_len], device=device, dtype=torch.int32)
739
- flash_attn_varlen_func(
733
+ return flash_attn_varlen_func(
734
q=query,
735
k=k_full,
736
v=v_full,
@@ -746,9 +740,7 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
740
max_seqlen_k=seq_len,
741
softmax_scale=self.scale,
742
causal=True,
749
- out=output,
743
)
751
- return output
744
else:
745
# SDPA fallback: expand KV for GQA, build causal mask
746
q_t = query.transpose(0, 1).unsqueeze(0) # (1, Hq, q_len, D)
vllm/v1/attention/ops/triton_turboquant_decode.py
+6
-6
@@ -143,12 +143,12 @@ def _tq_decode_stage1(
143
Block_table_ptr + bt_base + page_idx,
144
mask=kv_mask,
145
other=0,
146
- )
146
+ ).to(tl.int64)
147
148
slot_bases = (
149
block_nums * stride_cache_block
150
- + page_off * stride_cache_pos
151
- + kv_head * stride_cache_head
150
+ + page_off.to(tl.int64) * stride_cache_pos
151
+ + tl.cast(kv_head, tl.int64) * stride_cache_head
152
)
153
154
# ============================================================
@@ -356,11 +356,11 @@ def _tq_full_dequant_kv(
356
357
page_idx = pos // BLOCK_SIZE
358
page_off = pos % BLOCK_SIZE
359
- block_num = tl.load(Block_table_ptr + bid * stride_bt_b + page_idx)
359
+ block_num = tl.load(Block_table_ptr + bid * stride_bt_b + page_idx).to(tl.int64)
360
slot_base = (
361
block_num * stride_cache_block
362
- + page_off * stride_cache_pos
363
- + hid * stride_cache_head
362
+ + tl.cast(page_off, tl.int64) * stride_cache_pos
363
+ + tl.cast(hid, tl.int64) * stride_cache_head
364
)
365
366
d_offs = tl.arange(0, BLOCK_D)
vllm/v1/attention/ops/triton_turboquant_store.py
+12
-6
@@ -174,10 +174,13 @@ def _tq_fused_store_fp8(
174
slot = tl.load(Slot_mapping_ptr + token_idx)
175
if slot < 0:
176
return
177
- blk = slot // BLOCK_SIZE
178
- off = slot % BLOCK_SIZE
177
+ blk = (slot // BLOCK_SIZE).to(tl.int64)
178
+ off = (slot % BLOCK_SIZE).to(tl.int64)
179
+ head_idx_i64 = tl.cast(head_idx, tl.int64)
180
slot_base = (
180
- blk * stride_cache_block + off * stride_cache_pos + head_idx * stride_cache_head
181
+ blk * stride_cache_block
182
+ + off * stride_cache_pos
183
+ + head_idx_i64 * stride_cache_head
184
)
185
186
base = pid * D
@@ -259,10 +262,13 @@ def _tq_fused_store_mse(
262
slot = tl.load(Slot_mapping_ptr + token_idx)
263
if slot < 0:
264
return
262
- blk = slot // BLOCK_SIZE
263
- off = slot % BLOCK_SIZE
265
+ blk = (slot // BLOCK_SIZE).to(tl.int64)
266
+ off = (slot % BLOCK_SIZE).to(tl.int64)
267
+ head_idx_i64 = tl.cast(head_idx, tl.int64)
268
slot_base = (
265
- blk * stride_cache_block + off * stride_cache_pos + head_idx * stride_cache_head
269
+ blk * stride_cache_block
270
+ + off * stride_cache_pos
271
+ + head_idx_i64 * stride_cache_head
272
)
273
274
base = pid * D