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, 16:00 UTC
[Bugfix] Restrict TRTLLM attention to SM100, fixing GB300 (SM103) hang (#38730)
Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com>
Stefano Castagnetta committed
Apr 1, 2026 at 21:08 UTC
6183cae1bd8db73ee579f6424382975b17d02fb0
3 files changed
+14
-10
docs/design/attention_backends.md
+1
-1
@@ -167,7 +167,7 @@ Priority is **1 = highest** (tried first).
167
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ |
168
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A |
169
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
170
-| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
170
+| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.0 |
171
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
172
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
173
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
tools/pre_commit/generate_attention_backend_docs.py
+9
-5
@@ -235,10 +235,11 @@ def _resolve_import_to_file(
235
236
237
def _find_cc_in_function(tree: ast.AST, func_name: str) -> str | None:
238
- """Find a compute capability from is_device_capability_family() calls in a function.
238
+ """Find a compute capability from is_device_capability*() calls in a function.
239
240
- Looks for the pattern: current_platform.is_device_capability_family(N)
241
- and converts N (e.g. 100) to a CC string (e.g. "10.x").
240
+ Handles two patterns:
241
+ - is_device_capability_family(N): "M.x" (e.g. 100 -> "10.x")
242
+ - is_device_capability(N): "M.m" (e.g. 100 -> "10.0")
243
"""
244
for node in ast.walk(tree):
245
if not isinstance(node, ast.FunctionDef) or node.name != func_name:
@@ -247,12 +248,15 @@ def _find_cc_in_function(tree: ast.AST, func_name: str) -> str | None:
248
if (
249
isinstance(n, ast.Call)
250
and isinstance(n.func, ast.Attribute)
250
- and n.func.attr == "is_device_capability_family"
251
and n.args
252
and isinstance(n.args[0], ast.Constant)
253
and isinstance(n.args[0].value, int)
254
):
255
- return f"{n.args[0].value // 10}.x"
255
+ val = n.args[0].value
256
+ if n.func.attr == "is_device_capability_family":
257
+ return f"{val // 10}.x"
258
+ elif n.func.attr == "is_device_capability":
259
+ return f"{val // 10}.{val % 10}"
260
return None
261
262
vllm/utils/flashinfer.py
+4
-4
@@ -290,10 +290,10 @@ def supports_trtllm_attention() -> bool:
290
if envs.VLLM_BATCH_INVARIANT:
291
return False
292
293
- # Requires SM100 and NVIDIA artifactory to be accessible to download cubins
294
- return (
295
- current_platform.is_device_capability_family(100) and has_nvidia_artifactory()
296
- )
293
+ # TRTLLM attention is currently only validated on SM100 (CC 10.0).
294
+ # SM103 (GB300) hangs with FlashInfer >= 0.6.7.
295
+ # See: https://github.com/flashinfer-ai/flashinfer/issues/2939
296
+ return current_platform.is_device_capability(100) and has_nvidia_artifactory()
297
298
299
def force_use_trtllm_attention() -> bool | None: