[vasunvidia]

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Confidence Score: 0/5

[Medium impact] The changes affect eligible FP8 quantization calls and a specific fused-attention path.

This PR is not ready to merge because the six previous failures remain, and concurrent graph replay can also corrupt quantization results.

Findings

1.  Graphs overwrite each other's scale ▶
2.  Attention fallback raises ▶
3.  KV cache reads empty placeholders ▶
4.  Some head counts cannot compile ▶
5.  Inference mode raises ▶
6.  Streams read unfinished tables ▶
7.  Ordinary attention now needs Triton ▶

Summary

This PR adds fused RoPE and FP8 current-scaling quantization.

• Eligible attention calls can rotate and quantize Q/K/V together.
• Large BF16 tensors can use a two-pass Triton FP8 quantizer.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Combined QKV] --> B{Defer RoPE?}
  B -->|No| C[Regular RoPE]
  B -->|Yes| D[QKV placeholders]
  D --> E{Compatible FP8 attention?}
  E -->|Yes| F[Fused RoPE and FP8 quantization]
  E -->|No| G[Fill placeholders with regular RoPE]
  C --> H[Attention]
  G --> H
  F --> H

Loading

_{Reviews (2) · Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." · Reviewed by Greptile}

[greptile-apps]

Attention fallback raises

With checkpoint_core_attention=True, or FP8 forward with NVTE_FP8_DPA_BWD=0, DotProductAttention first calls materialize(). That sets state.consumed=True but leaves the state attached to q. The fused backend then calls quantize_deferred(), which raises “already consumed” instead of quantizing the filled tensors. Return None for an already materialized state so the regular quantization path can run.

Knowledge Base Used: PyTorch attention execution

[greptile-apps]

KV cache reads empty placeholders

The new materialize() call runs after inference_params.step(). For an sbhd prefill using deferred RoPE, that earlier step already copies the uninitialized K/V placeholders into the cache and replaces key_layer and value_layer with cache tensors. The fallback then copies sbhd RoPE results into the bshd cache, which can raise for mismatched shapes or overwrite cached data incorrectly. Materialize the original placeholders before any KV-cache update.

Knowledge Base Used: PyTorch attention execution

[greptile-apps]

Some head counts cannot compile

G and HQ are passed directly to tl.arange(), which requires power-of-two range sizes. The deferred input checks allow other head counts, such as 12 query heads across 3 groups (G=3, HQ=4). These inputs reach the kernel and fail to compile instead of using the promised fallback. Pad these ranges and mask their loads and stores, or reject unsupported counts before returning placeholders.

[greptile-apps]

Inference mode raises

Accessing freqs._version raises for tensors created under torch.inference_mode(). A user who creates rotary embeddings inside that context and runs full-sequence FP8 attention without InferenceParams reaches this function through the deferred path. Even the first cache miss raises when storing the entry. Handle inference tensors without reading their version counter, for example by computing their tables without this cache.

[greptile-apps]

Streams read unfinished tables

The table cache is shared across CUDA streams, but a cache hit does not wait for the stream that computed cos and sin. If two attention calls share an already-ready frequency tensor on different streams, the second call can reuse the entry while the first stream is still writing its tables. It then rotates Q/K using unfinished values. Cache tables per stream or record an event and wait before another stream uses them.

[greptile-apps]

Ordinary attention now needs Triton

quantize_deferred() imports the Triton wrapper before checking whether q has deferred state. If Triton is absent, every ordinary FP8 fused-attention call now fails here, even when deferred RoPE is not requested. The new automatic large-tensor quantization path in float8_tensor.py also imports Triton without a fallback. The declared PyTorch requirements do not include Triton, and the existing fused MLA code treats it as optional. Move the deferred import below the state check and retain the native path when Triton is unavailable.

[greptile-apps]

Graphs overwrite each other's scale

Capturing two CUDA graphs on the same stream makes both reuse the same _WORKSPACES entry. If those graphs later replay on different streams, their kernels update the same maximum, counter, and scale at once. One graph can reset the counter or overwrite the scale while the other still needs it, producing incorrect FP8 values.

Allocate separate scratch space for each captured graph instead of sharing it by capture stream.